diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/BytesUtils.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/BytesUtils.kt new file mode 100644 index 000000000..9994ba56c --- /dev/null +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/BytesUtils.kt @@ -0,0 +1,5 @@ +package org.jetbrains.kotlinx.dataframe.impl.io + +import java.util.Base64 + +internal fun ByteArray.toBase64(): String = Base64.getEncoder().encodeToString(this) diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/compression.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/compression.kt new file mode 100644 index 000000000..8d95b811e --- /dev/null +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/compression.kt @@ -0,0 +1,11 @@ +package org.jetbrains.kotlinx.dataframe.impl.io + +import java.io.ByteArrayOutputStream +import java.util.zip.GZIPOutputStream + +internal fun ByteArray.encodeGzip(): ByteArray { + val bos = ByteArrayOutputStream() + GZIPOutputStream(bos).use { it.write(this) } + + return bos.toByteArray() +} diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/image.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/image.kt new file mode 100644 index 000000000..853353230 --- /dev/null +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/image.kt @@ -0,0 +1,60 @@ +package org.jetbrains.kotlinx.dataframe.impl.io + +import java.awt.Graphics2D +import java.awt.RenderingHints +import java.awt.image.BufferedImage +import java.awt.image.ImageObserver +import java.io.ByteArrayOutputStream +import javax.imageio.ImageIO +import kotlin.math.max +import kotlin.math.min + +internal fun BufferedImage.resizeKeepingAspectRatio( + maxSize: Int, + resultImageType: Int = BufferedImage.TYPE_INT_ARGB, + interpolation: Any = RenderingHints.VALUE_INTERPOLATION_NEAREST_NEIGHBOR, + renderingQuality: Any = RenderingHints.VALUE_RENDER_QUALITY, + antialiasing: Any = RenderingHints.VALUE_ANTIALIAS_ON, + observer: ImageObserver? = null +): BufferedImage { + val aspectRatio = width.toDouble() / height.toDouble() + val size = min(maxSize, max(width, height)) + + val (nWidth, nHeight) = if (width > height) { + Pair(size, (size / aspectRatio).toInt()) + } else { + Pair((size * aspectRatio).toInt(), size) + } + + return resize(nWidth, nHeight, resultImageType, interpolation, renderingQuality, antialiasing, observer) +} + +internal fun BufferedImage.resize( + width: Int, + height: Int, + resultImageType: Int = BufferedImage.TYPE_INT_ARGB, + interpolation: Any = RenderingHints.VALUE_INTERPOLATION_NEAREST_NEIGHBOR, + renderingQuality: Any = RenderingHints.VALUE_RENDER_QUALITY, + antialiasing: Any = RenderingHints.VALUE_ANTIALIAS_ON, + observer: ImageObserver? = null +): BufferedImage { + val resized = BufferedImage(width, height, resultImageType) + val g: Graphics2D = resized.createGraphics() + + g.setRenderingHint(RenderingHints.KEY_INTERPOLATION, interpolation) + g.setRenderingHint(RenderingHints.KEY_RENDERING, renderingQuality) + g.setRenderingHint(RenderingHints.KEY_ANTIALIASING, antialiasing) + + g.drawImage(this, 0, 0, width, height, observer) + g.dispose() + + return resized +} + +internal const val DEFAULT_IMG_FORMAT = "png" + +internal fun BufferedImage.toByteArray(format: String = DEFAULT_IMG_FORMAT): ByteArray = + ByteArrayOutputStream().use { bos -> + ImageIO.write(this, format, bos) + bos.toByteArray() + } diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/writeJson.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/writeJson.kt index cc494c507..4c8290a84 100644 --- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/writeJson.kt +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/writeJson.kt @@ -23,11 +23,14 @@ import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.METADATA import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.NCOL import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.NROW import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.VERSION +import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions import org.jetbrains.kotlinx.dataframe.io.arrayColumnName import org.jetbrains.kotlinx.dataframe.io.valueColumnName import org.jetbrains.kotlinx.dataframe.ncol import org.jetbrains.kotlinx.dataframe.nrow import org.jetbrains.kotlinx.dataframe.typeClass +import java.awt.image.BufferedImage +import java.io.IOException internal fun KlaxonJson.encodeRow(frame: ColumnsContainer<*>, index: Int): JsonObject? { val values = frame.columns().map { col -> @@ -57,18 +60,19 @@ internal const val SERIALIZATION_VERSION = "2.0.0" internal fun KlaxonJson.encodeRowWithMetadata( frame: ColumnsContainer<*>, index: Int, - rowLimit: Int? = null + rowLimit: Int? = null, + imageEncodingOptions: Base64ImageEncodingOptions? = null ): JsonObject? { val values = frame.columns().map { col -> when (col) { is ColumnGroup<*> -> obj( - DATA to encodeRowWithMetadata(col, index, rowLimit), + DATA to encodeRowWithMetadata(col, index, rowLimit, imageEncodingOptions), METADATA to obj(KIND to ColumnKind.Group.toString()) ) is FrameColumn<*> -> { - val data = if (rowLimit == null) encodeFrameWithMetadata(col[index]) - else encodeFrameWithMetadata(col[index].take(rowLimit), rowLimit) + val data = if (rowLimit == null) encodeFrameWithMetadata(col[index], null, imageEncodingOptions) + else encodeFrameWithMetadata(col[index].take(rowLimit), rowLimit, imageEncodingOptions) obj( DATA to data, METADATA to obj( @@ -79,7 +83,7 @@ internal fun KlaxonJson.encodeRowWithMetadata( ) } - else -> encodeValue(col, index) + else -> encodeValue(col, index, imageEncodingOptions) }.let { col.name to it } } if (values.isEmpty()) return null @@ -89,7 +93,11 @@ internal fun KlaxonJson.encodeRowWithMetadata( private val valueTypes = setOf(Boolean::class, Double::class, Int::class, Float::class, Long::class, Short::class, Byte::class) -internal fun KlaxonJson.encodeValue(col: AnyCol, index: Int): Any? = when { +internal fun KlaxonJson.encodeValue( + col: AnyCol, + index: Int, + imageEncodingOptions: Base64ImageEncodingOptions? = null +): Any? = when { col.isList() -> col[index]?.let { list -> val values = (list as List<*>).map { when (it) { @@ -101,6 +109,7 @@ internal fun KlaxonJson.encodeValue(col: AnyCol, index: Int): Any? = when { } array(values) } ?: array() + col.typeClass in valueTypes -> { val v = col[index] if ((v is Double && v.isNaN()) || (v is Float && v.isNaN())) { @@ -108,10 +117,42 @@ internal fun KlaxonJson.encodeValue(col: AnyCol, index: Int): Any? = when { } else v } + col.typeClass == BufferedImage::class && imageEncodingOptions != null -> + col[index]?.let { image -> + encodeBufferedImageAsBase64(image as BufferedImage, imageEncodingOptions) + } ?: "" + else -> col[index]?.toString() } -internal fun KlaxonJson.encodeFrameWithMetadata(frame: AnyFrame, rowLimit: Int? = null): JsonArray<*> { +private fun encodeBufferedImageAsBase64( + image: BufferedImage, + imageEncodingOptions: Base64ImageEncodingOptions = Base64ImageEncodingOptions() +): String? { + return try { + val preparedImage = if (imageEncodingOptions.isLimitSizeOn) { + image.resizeKeepingAspectRatio(imageEncodingOptions.imageSizeLimit) + } else { + image + } + + val bytes = if (imageEncodingOptions.isGzipOn) { + preparedImage.toByteArray().encodeGzip() + } else { + preparedImage.toByteArray() + } + + bytes.toBase64() + } catch (e: IOException) { + null + } +} + +internal fun KlaxonJson.encodeFrameWithMetadata( + frame: AnyFrame, + rowLimit: Int? = null, + imageEncodingOptions: Base64ImageEncodingOptions? = null +): JsonArray<*> { val valueColumn = frame.extractValueColumn() val arrayColumn = frame.extractArrayColumn() @@ -122,9 +163,13 @@ internal fun KlaxonJson.encodeFrameWithMetadata(frame: AnyFrame, rowLimit: Int? ?.get(rowIndex) ?: arrayColumn?.get(rowIndex) ?.let { - if (arraysAreFrames) encodeFrameWithMetadata(it as AnyFrame, rowLimit) else null + if (arraysAreFrames) encodeFrameWithMetadata( + it as AnyFrame, + rowLimit, + imageEncodingOptions + ) else null } - ?: encodeRowWithMetadata(frame, rowIndex, rowLimit) + ?: encodeRowWithMetadata(frame, rowIndex, rowLimit, imageEncodingOptions) } return array(data) @@ -206,6 +251,7 @@ internal fun KlaxonJson.encodeDataFrameWithMetadata( frame: AnyFrame, rowLimit: Int, nestedRowLimit: Int? = null, + imageEncodingOptions: Base64ImageEncodingOptions? = null ): JsonObject { return obj( VERSION to SERIALIZATION_VERSION, @@ -216,7 +262,8 @@ internal fun KlaxonJson.encodeDataFrameWithMetadata( ), KOTLIN_DATAFRAME to encodeFrameWithMetadata( frame.take(rowLimit), - rowLimit = nestedRowLimit + rowLimit = nestedRowLimit, + imageEncodingOptions ), ) } diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/html.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/html.kt index 763c904dc..60eac3042 100644 --- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/html.kt +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/html.kt @@ -21,6 +21,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath import org.jetbrains.kotlinx.dataframe.columns.FrameColumn import org.jetbrains.kotlinx.dataframe.impl.DataFrameSize import org.jetbrains.kotlinx.dataframe.impl.columns.addPath +import org.jetbrains.kotlinx.dataframe.impl.io.resizeKeepingAspectRatio import org.jetbrains.kotlinx.dataframe.impl.renderType import org.jetbrains.kotlinx.dataframe.impl.scale import org.jetbrains.kotlinx.dataframe.impl.truncate @@ -31,6 +32,7 @@ import org.jetbrains.kotlinx.dataframe.name import org.jetbrains.kotlinx.dataframe.nrow import org.jetbrains.kotlinx.dataframe.size import java.awt.Desktop +import java.awt.image.BufferedImage import java.io.File import java.io.InputStreamReader import java.net.URL @@ -152,7 +154,8 @@ internal fun AnyFrame.toHtmlData( DataFrameReference(id, value.size) } } else { - val html = formatter.format(value, cellRenderer, renderConfig) + val html = + formatter.format(downsizeBufferedImageIfNeeded(value, renderConfig), cellRenderer, renderConfig) val style = renderConfig.cellFormatter?.invoke(FormattingDSL, it, col)?.attributes()?.ifEmpty { null } ?.joinToString(";") { "${it.first}:${it.second}" } HtmlContent(html, style) @@ -180,6 +183,26 @@ internal fun AnyFrame.toHtmlData( return DataFrameHtmlData(style = "", body = body, script = script) } +private const val DEFAULT_HTML_IMG_SIZE = 100 + +/** + * This method resizes a BufferedImage if necessary, according to the provided DisplayConfiguration. + * It is essential to prevent potential memory problems when serializing HTML data for display in the Kotlin Notebook plugin. + * + * @param value The input value to be checked and possibly downsized. + * @param renderConfig The DisplayConfiguration to determine if downsizing is needed. + * @return The downsized BufferedImage if value is a BufferedImage and downsizing is enabled in the DisplayConfiguration, + * otherwise returns the input value unchanged. + */ +private fun downsizeBufferedImageIfNeeded(value: Any?, renderConfig: DisplayConfiguration): Any? = + when { + value is BufferedImage && renderConfig.downsizeBufferedImage -> { + value.resizeKeepingAspectRatio(DEFAULT_HTML_IMG_SIZE) + } + + else -> value + } + /** * Renders [this] [DataFrame] as static HTML (meaning no JS is used). * CSS rendering is enabled by default but can be turned off using [includeCss] @@ -568,6 +591,7 @@ public data class DisplayConfiguration( internal val localTesting: Boolean = flagFromEnv("KOTLIN_DATAFRAME_LOCAL_TESTING"), var useDarkColorScheme: Boolean = false, var enableFallbackStaticTables: Boolean = true, + var downsizeBufferedImage: Boolean = true ) { public companion object { public val DEFAULT: DisplayConfiguration = DisplayConfiguration() diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt index e648458f9..295a40413 100644 --- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt @@ -276,6 +276,7 @@ public fun AnyFrame.toJson(prettyPrint: Boolean = false, canonical: Boolean = fa * Applied for each frame column recursively * @param prettyPrint Specifies whether the output JSON should be formatted with indentation and line breaks. * @param canonical Specifies whether the output JSON should be in a canonical form. + * @param imageEncodingOptions The options for encoding images. The default is null, which indicates that the image is not encoded as Base64. * * @return The DataFrame converted to a JSON string with metadata. */ @@ -283,13 +284,39 @@ public fun AnyFrame.toJsonWithMetadata( rowLimit: Int, nestedRowLimit: Int? = null, prettyPrint: Boolean = false, - canonical: Boolean = false + canonical: Boolean = false, + imageEncodingOptions: Base64ImageEncodingOptions? = null ): String { return json { - encodeDataFrameWithMetadata(this@toJsonWithMetadata, rowLimit, nestedRowLimit) + encodeDataFrameWithMetadata(this@toJsonWithMetadata, rowLimit, nestedRowLimit, imageEncodingOptions) }.toJsonString(prettyPrint, canonical) } +internal const val DEFAULT_IMG_SIZE = 600 + +/** + * Class representing the options for encoding images. + * + * @property imageSizeLimit The maximum size to which images should be resized. Defaults to the value of DEFAULT_IMG_SIZE. + * @property options Bitwise-OR of the [GZIP_ON] and [LIMIT_SIZE_ON] constants. Defaults to [GZIP_ON] or [LIMIT_SIZE_ON]. + */ +public class Base64ImageEncodingOptions( + public val imageSizeLimit: Int = DEFAULT_IMG_SIZE, + private val options: Int = GZIP_ON or LIMIT_SIZE_ON +) { + public val isGzipOn: Boolean + get() = options and GZIP_ON == GZIP_ON + + public val isLimitSizeOn: Boolean + get() = options and LIMIT_SIZE_ON == LIMIT_SIZE_ON + + public companion object { + public const val ALL_OFF: Int = 0 + public const val GZIP_ON: Int = 1 // 2^0 + public const val LIMIT_SIZE_ON: Int = 2 // 2^1 + } +} + public fun AnyRow.toJson(prettyPrint: Boolean = false, canonical: Boolean = false): String { return json { encodeRow(df(), index()) diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/JupyterHtmlRenderer.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/JupyterHtmlRenderer.kt index 536470fa8..a5a4dc249 100644 --- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/JupyterHtmlRenderer.kt +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/JupyterHtmlRenderer.kt @@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.jupyter import com.beust.klaxon.json import org.jetbrains.kotlinx.dataframe.api.take import org.jetbrains.kotlinx.dataframe.impl.io.encodeFrame +import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions import org.jetbrains.kotlinx.dataframe.io.DataFrameHtmlData import org.jetbrains.kotlinx.dataframe.io.DisplayConfiguration import org.jetbrains.kotlinx.dataframe.io.toHTML @@ -23,6 +24,7 @@ import org.jetbrains.kotlinx.jupyter.api.renderHtmlAsIFrameIfNeeded /** Starting from this version, dataframe integration will respond with additional data for rendering in Kotlin Notebooks plugin. */ private const val MIN_KERNEL_VERSION_FOR_NEW_TABLES_UI = "0.11.0.311" private const val MIN_IDE_VERSION_SUPPORT_JSON_WITH_METADATA = 241 +private const val MIN_IDE_VERSION_SUPPORT_IMAGE_VIEWER = 242 internal class JupyterHtmlRenderer( val display: DisplayConfiguration, @@ -63,8 +65,8 @@ internal inline fun JupyterHtmlRenderer.render( if (notebook.kernelVersion >= KotlinKernelVersion.from(MIN_KERNEL_VERSION_FOR_NEW_TABLES_UI)!!) { val ideBuildNumber = KotlinNotebookPluginUtils.getKotlinNotebookIDEBuildNumber() - val jsonEncodedDf = - if (ideBuildNumber == null || ideBuildNumber.majorVersion < MIN_IDE_VERSION_SUPPORT_JSON_WITH_METADATA) { + val jsonEncodedDf = when { + !ideBuildNumber.supportsDynamicNestedTables() -> { json { obj( "nrow" to df.size.nrow, @@ -73,15 +75,32 @@ internal inline fun JupyterHtmlRenderer.render( "kotlin_dataframe" to encodeFrame(df.take(limit)), ) }.toJsonString() - } else { - df.toJsonWithMetadata(limit, reifiedDisplayConfiguration.rowsLimit) } + + else -> { + val imageEncodingOptions = + if (ideBuildNumber.supportsImageViewer()) Base64ImageEncodingOptions() else null + + df.toJsonWithMetadata( + limit, + reifiedDisplayConfiguration.rowsLimit, + imageEncodingOptions = imageEncodingOptions + ) + } + } + notebook.renderAsIFrameAsNeeded(html, staticHtml, jsonEncodedDf) } else { notebook.renderHtmlAsIFrameIfNeeded(html) } } +private fun KotlinNotebookPluginUtils.IdeBuildNumber?.supportsDynamicNestedTables() = + this != null && majorVersion >= MIN_IDE_VERSION_SUPPORT_JSON_WITH_METADATA + +private fun KotlinNotebookPluginUtils.IdeBuildNumber?.supportsImageViewer() = + this != null && majorVersion >= MIN_IDE_VERSION_SUPPORT_IMAGE_VIEWER + internal fun Notebook.renderAsIFrameAsNeeded( data: HtmlData, staticData: HtmlData, diff --git a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/Utils.kt b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/Utils.kt index c071955e0..dae5f687b 100644 --- a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/Utils.kt +++ b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/Utils.kt @@ -1,5 +1,7 @@ package org.jetbrains.kotlinx.dataframe +import com.beust.klaxon.JsonObject +import com.beust.klaxon.Parser import org.jetbrains.kotlinx.dataframe.api.print import org.jetbrains.kotlinx.dataframe.api.schema import org.jetbrains.kotlinx.dataframe.io.renderToString @@ -24,3 +26,8 @@ fun > T.alsoDebug(println: String? = null, rowsLimit: Int = 20) print(borders = true, title = true, columnTypes = true, valueLimit = -1, rowsLimit = rowsLimit) schema().print() } + +fun parseJsonStr(jsonStr: String): JsonObject { + val parser = Parser.default() + return parser.parse(StringBuilder(jsonStr)) as JsonObject +} diff --git a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/ImageSerializationTests.kt b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/ImageSerializationTests.kt new file mode 100644 index 000000000..cae6e759c --- /dev/null +++ b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/ImageSerializationTests.kt @@ -0,0 +1,174 @@ +package org.jetbrains.kotlinx.dataframe.io + +import com.beust.klaxon.JsonArray +import com.beust.klaxon.JsonObject +import io.kotest.matchers.shouldBe +import io.kotest.matchers.string.shouldContain +import org.jetbrains.kotlinx.dataframe.api.dataFrameOf +import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.KOTLIN_DATAFRAME +import org.jetbrains.kotlinx.dataframe.impl.io.resizeKeepingAspectRatio +import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions.Companion.ALL_OFF +import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions.Companion.GZIP_ON +import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions.Companion.LIMIT_SIZE_ON +import org.jetbrains.kotlinx.dataframe.parseJsonStr +import org.jetbrains.kotlinx.dataframe.testResource +import org.junit.Test +import org.junit.runner.RunWith +import org.junit.runners.Parameterized +import java.awt.image.BufferedImage +import java.io.ByteArrayInputStream +import java.io.ByteArrayOutputStream +import java.io.File +import java.util.* +import java.util.zip.GZIPInputStream +import javax.imageio.ImageIO + +@RunWith(Parameterized::class) +class ImageSerializationTests(private val encodingOptions: Base64ImageEncodingOptions?) { + @Test + fun `serialize images as base64`() { + val images = readImagesFromResources() + val json = encodeImagesAsJson(images, encodingOptions) + + if (encodingOptions == DISABLED) { + checkImagesEncodedAsToString(json, images.size) + return + } + + val decodedImages = decodeImagesFromJson(json, images.size, encodingOptions!!) + + for ((decodedImage, original) in decodedImages.zip(images)) { + val expectedImage = resizeIfNeeded(original, encodingOptions) + isImagesIdentical(decodedImage, expectedImage, 2) shouldBe true + } + } + + private fun readImagesFromResources(): List { + val dir = File(testResource("imgs").path) + + return dir.listFiles()?.map { file -> + try { + ImageIO.read(file) + } catch (ex: Exception) { + throw IllegalArgumentException("Error reading ${file.name}: ${ex.message}") + } + } ?: emptyList() + } + + private fun encodeImagesAsJson( + images: List, + encodingOptions: Base64ImageEncodingOptions? + ): JsonObject { + val df = dataFrameOf(listOf("imgs"), images) + val jsonStr = df.toJsonWithMetadata(20, nestedRowLimit = 20, imageEncodingOptions = encodingOptions) + + return parseJsonStr(jsonStr) + } + + private fun checkImagesEncodedAsToString(json: JsonObject, numImgs: Int) { + for (i in 0..)[i] as JsonObject + val img = row["imgs"] as String + + img shouldContain "BufferedImage" + } + } + + private fun decodeImagesFromJson( + json: JsonObject, + imgsNum: Int, + encodingOptions: Base64ImageEncodingOptions + ): List { + val result = mutableListOf() + for (i in 0..)[i] as JsonObject + val imgString = row["imgs"] as String + + val bytes = decodeBase64Image(imgString, encodingOptions) + val decodedImage = createImageFromBytes(bytes) + + result.add(decodedImage) + } + + return result + } + + private fun decodeBase64Image(imgString: String, encodingOptions: Base64ImageEncodingOptions): ByteArray = + when { + encodingOptions.isGzipOn -> decompressGzip(Base64.getDecoder().decode(imgString)) + else -> Base64.getDecoder().decode(imgString) + } + + private fun decompressGzip(input: ByteArray): ByteArray { + return ByteArrayOutputStream().use { byteArrayOutputStream -> + GZIPInputStream(input.inputStream()).use { inputStream -> + inputStream.copyTo(byteArrayOutputStream) + } + byteArrayOutputStream.toByteArray() + } + } + + private fun resizeIfNeeded(image: BufferedImage, encodingOptions: Base64ImageEncodingOptions): BufferedImage = + when { + !encodingOptions.isLimitSizeOn -> image + else -> image.resizeKeepingAspectRatio(encodingOptions.imageSizeLimit) + } + + private fun createImageFromBytes(bytes: ByteArray): BufferedImage { + val bais = ByteArrayInputStream(bytes) + return ImageIO.read(bais) + } + + private fun isImagesIdentical(img1: BufferedImage, img2: BufferedImage, allowedDelta: Int): Boolean { + // First check dimensions + if (img1.width != img2.width || img1.height != img2.height) { + return false + } + + // Then check each pixel + for (y in 0 until img1.height) { + for (x in 0 until img1.width) { + val rgb1 = img1.getRGB(x, y) + val rgb2 = img2.getRGB(x, y) + + val r1 = (rgb1 shr 16) and 0xFF + val g1 = (rgb1 shr 8) and 0xFF + val b1 = rgb1 and 0xFF + + val r2 = (rgb2 shr 16) and 0xFF + val g2 = (rgb2 shr 8) and 0xFF + val b2 = rgb2 and 0xFF + + val diff = kotlin.math.abs(r1 - r2) + kotlin.math.abs(g1 - g2) + kotlin.math.abs(b1 - b2) + + // If the difference in color components exceed our allowance return false + if (diff > allowedDelta) { + return false + } + } + } + + // If no exceeding difference was found, the images are identical within our allowedDelta + return true + } + + companion object { + private val DEFAULT = Base64ImageEncodingOptions() + private val GZIP_ON_RESIZE_OFF = Base64ImageEncodingOptions(options = GZIP_ON) + private val GZIP_OFF_RESIZE_OFF = Base64ImageEncodingOptions(options = ALL_OFF) + private val GZIP_ON_RESIZE_TO_700 = Base64ImageEncodingOptions(imageSizeLimit = 700, options = GZIP_ON or LIMIT_SIZE_ON) + private val DISABLED = null + + @JvmStatic + @Parameterized.Parameters + fun imageEncodingOptionsToTest(): Collection { + return listOf( + DEFAULT, + GZIP_ON_RESIZE_OFF, + GZIP_OFF_RESIZE_OFF, + GZIP_ON_RESIZE_TO_700, + null, + ) + } + } +} diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/BytesUtils.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/BytesUtils.kt new file mode 100644 index 000000000..9994ba56c --- /dev/null +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/BytesUtils.kt @@ -0,0 +1,5 @@ +package org.jetbrains.kotlinx.dataframe.impl.io + +import java.util.Base64 + +internal fun ByteArray.toBase64(): String = Base64.getEncoder().encodeToString(this) diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/compression.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/compression.kt new file mode 100644 index 000000000..8d95b811e --- /dev/null +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/compression.kt @@ -0,0 +1,11 @@ +package org.jetbrains.kotlinx.dataframe.impl.io + +import java.io.ByteArrayOutputStream +import java.util.zip.GZIPOutputStream + +internal fun ByteArray.encodeGzip(): ByteArray { + val bos = ByteArrayOutputStream() + GZIPOutputStream(bos).use { it.write(this) } + + return bos.toByteArray() +} diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/image.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/image.kt new file mode 100644 index 000000000..853353230 --- /dev/null +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/image.kt @@ -0,0 +1,60 @@ +package org.jetbrains.kotlinx.dataframe.impl.io + +import java.awt.Graphics2D +import java.awt.RenderingHints +import java.awt.image.BufferedImage +import java.awt.image.ImageObserver +import java.io.ByteArrayOutputStream +import javax.imageio.ImageIO +import kotlin.math.max +import kotlin.math.min + +internal fun BufferedImage.resizeKeepingAspectRatio( + maxSize: Int, + resultImageType: Int = BufferedImage.TYPE_INT_ARGB, + interpolation: Any = RenderingHints.VALUE_INTERPOLATION_NEAREST_NEIGHBOR, + renderingQuality: Any = RenderingHints.VALUE_RENDER_QUALITY, + antialiasing: Any = RenderingHints.VALUE_ANTIALIAS_ON, + observer: ImageObserver? = null +): BufferedImage { + val aspectRatio = width.toDouble() / height.toDouble() + val size = min(maxSize, max(width, height)) + + val (nWidth, nHeight) = if (width > height) { + Pair(size, (size / aspectRatio).toInt()) + } else { + Pair((size * aspectRatio).toInt(), size) + } + + return resize(nWidth, nHeight, resultImageType, interpolation, renderingQuality, antialiasing, observer) +} + +internal fun BufferedImage.resize( + width: Int, + height: Int, + resultImageType: Int = BufferedImage.TYPE_INT_ARGB, + interpolation: Any = RenderingHints.VALUE_INTERPOLATION_NEAREST_NEIGHBOR, + renderingQuality: Any = RenderingHints.VALUE_RENDER_QUALITY, + antialiasing: Any = RenderingHints.VALUE_ANTIALIAS_ON, + observer: ImageObserver? = null +): BufferedImage { + val resized = BufferedImage(width, height, resultImageType) + val g: Graphics2D = resized.createGraphics() + + g.setRenderingHint(RenderingHints.KEY_INTERPOLATION, interpolation) + g.setRenderingHint(RenderingHints.KEY_RENDERING, renderingQuality) + g.setRenderingHint(RenderingHints.KEY_ANTIALIASING, antialiasing) + + g.drawImage(this, 0, 0, width, height, observer) + g.dispose() + + return resized +} + +internal const val DEFAULT_IMG_FORMAT = "png" + +internal fun BufferedImage.toByteArray(format: String = DEFAULT_IMG_FORMAT): ByteArray = + ByteArrayOutputStream().use { bos -> + ImageIO.write(this, format, bos) + bos.toByteArray() + } diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/writeJson.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/writeJson.kt index cc494c507..4c8290a84 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/writeJson.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/writeJson.kt @@ -23,11 +23,14 @@ import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.METADATA import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.NCOL import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.NROW import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.VERSION +import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions import org.jetbrains.kotlinx.dataframe.io.arrayColumnName import org.jetbrains.kotlinx.dataframe.io.valueColumnName import org.jetbrains.kotlinx.dataframe.ncol import org.jetbrains.kotlinx.dataframe.nrow import org.jetbrains.kotlinx.dataframe.typeClass +import java.awt.image.BufferedImage +import java.io.IOException internal fun KlaxonJson.encodeRow(frame: ColumnsContainer<*>, index: Int): JsonObject? { val values = frame.columns().map { col -> @@ -57,18 +60,19 @@ internal const val SERIALIZATION_VERSION = "2.0.0" internal fun KlaxonJson.encodeRowWithMetadata( frame: ColumnsContainer<*>, index: Int, - rowLimit: Int? = null + rowLimit: Int? = null, + imageEncodingOptions: Base64ImageEncodingOptions? = null ): JsonObject? { val values = frame.columns().map { col -> when (col) { is ColumnGroup<*> -> obj( - DATA to encodeRowWithMetadata(col, index, rowLimit), + DATA to encodeRowWithMetadata(col, index, rowLimit, imageEncodingOptions), METADATA to obj(KIND to ColumnKind.Group.toString()) ) is FrameColumn<*> -> { - val data = if (rowLimit == null) encodeFrameWithMetadata(col[index]) - else encodeFrameWithMetadata(col[index].take(rowLimit), rowLimit) + val data = if (rowLimit == null) encodeFrameWithMetadata(col[index], null, imageEncodingOptions) + else encodeFrameWithMetadata(col[index].take(rowLimit), rowLimit, imageEncodingOptions) obj( DATA to data, METADATA to obj( @@ -79,7 +83,7 @@ internal fun KlaxonJson.encodeRowWithMetadata( ) } - else -> encodeValue(col, index) + else -> encodeValue(col, index, imageEncodingOptions) }.let { col.name to it } } if (values.isEmpty()) return null @@ -89,7 +93,11 @@ internal fun KlaxonJson.encodeRowWithMetadata( private val valueTypes = setOf(Boolean::class, Double::class, Int::class, Float::class, Long::class, Short::class, Byte::class) -internal fun KlaxonJson.encodeValue(col: AnyCol, index: Int): Any? = when { +internal fun KlaxonJson.encodeValue( + col: AnyCol, + index: Int, + imageEncodingOptions: Base64ImageEncodingOptions? = null +): Any? = when { col.isList() -> col[index]?.let { list -> val values = (list as List<*>).map { when (it) { @@ -101,6 +109,7 @@ internal fun KlaxonJson.encodeValue(col: AnyCol, index: Int): Any? = when { } array(values) } ?: array() + col.typeClass in valueTypes -> { val v = col[index] if ((v is Double && v.isNaN()) || (v is Float && v.isNaN())) { @@ -108,10 +117,42 @@ internal fun KlaxonJson.encodeValue(col: AnyCol, index: Int): Any? = when { } else v } + col.typeClass == BufferedImage::class && imageEncodingOptions != null -> + col[index]?.let { image -> + encodeBufferedImageAsBase64(image as BufferedImage, imageEncodingOptions) + } ?: "" + else -> col[index]?.toString() } -internal fun KlaxonJson.encodeFrameWithMetadata(frame: AnyFrame, rowLimit: Int? = null): JsonArray<*> { +private fun encodeBufferedImageAsBase64( + image: BufferedImage, + imageEncodingOptions: Base64ImageEncodingOptions = Base64ImageEncodingOptions() +): String? { + return try { + val preparedImage = if (imageEncodingOptions.isLimitSizeOn) { + image.resizeKeepingAspectRatio(imageEncodingOptions.imageSizeLimit) + } else { + image + } + + val bytes = if (imageEncodingOptions.isGzipOn) { + preparedImage.toByteArray().encodeGzip() + } else { + preparedImage.toByteArray() + } + + bytes.toBase64() + } catch (e: IOException) { + null + } +} + +internal fun KlaxonJson.encodeFrameWithMetadata( + frame: AnyFrame, + rowLimit: Int? = null, + imageEncodingOptions: Base64ImageEncodingOptions? = null +): JsonArray<*> { val valueColumn = frame.extractValueColumn() val arrayColumn = frame.extractArrayColumn() @@ -122,9 +163,13 @@ internal fun KlaxonJson.encodeFrameWithMetadata(frame: AnyFrame, rowLimit: Int? ?.get(rowIndex) ?: arrayColumn?.get(rowIndex) ?.let { - if (arraysAreFrames) encodeFrameWithMetadata(it as AnyFrame, rowLimit) else null + if (arraysAreFrames) encodeFrameWithMetadata( + it as AnyFrame, + rowLimit, + imageEncodingOptions + ) else null } - ?: encodeRowWithMetadata(frame, rowIndex, rowLimit) + ?: encodeRowWithMetadata(frame, rowIndex, rowLimit, imageEncodingOptions) } return array(data) @@ -206,6 +251,7 @@ internal fun KlaxonJson.encodeDataFrameWithMetadata( frame: AnyFrame, rowLimit: Int, nestedRowLimit: Int? = null, + imageEncodingOptions: Base64ImageEncodingOptions? = null ): JsonObject { return obj( VERSION to SERIALIZATION_VERSION, @@ -216,7 +262,8 @@ internal fun KlaxonJson.encodeDataFrameWithMetadata( ), KOTLIN_DATAFRAME to encodeFrameWithMetadata( frame.take(rowLimit), - rowLimit = nestedRowLimit + rowLimit = nestedRowLimit, + imageEncodingOptions ), ) } diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/html.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/html.kt index 763c904dc..60eac3042 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/html.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/html.kt @@ -21,6 +21,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath import org.jetbrains.kotlinx.dataframe.columns.FrameColumn import org.jetbrains.kotlinx.dataframe.impl.DataFrameSize import org.jetbrains.kotlinx.dataframe.impl.columns.addPath +import org.jetbrains.kotlinx.dataframe.impl.io.resizeKeepingAspectRatio import org.jetbrains.kotlinx.dataframe.impl.renderType import org.jetbrains.kotlinx.dataframe.impl.scale import org.jetbrains.kotlinx.dataframe.impl.truncate @@ -31,6 +32,7 @@ import org.jetbrains.kotlinx.dataframe.name import org.jetbrains.kotlinx.dataframe.nrow import org.jetbrains.kotlinx.dataframe.size import java.awt.Desktop +import java.awt.image.BufferedImage import java.io.File import java.io.InputStreamReader import java.net.URL @@ -152,7 +154,8 @@ internal fun AnyFrame.toHtmlData( DataFrameReference(id, value.size) } } else { - val html = formatter.format(value, cellRenderer, renderConfig) + val html = + formatter.format(downsizeBufferedImageIfNeeded(value, renderConfig), cellRenderer, renderConfig) val style = renderConfig.cellFormatter?.invoke(FormattingDSL, it, col)?.attributes()?.ifEmpty { null } ?.joinToString(";") { "${it.first}:${it.second}" } HtmlContent(html, style) @@ -180,6 +183,26 @@ internal fun AnyFrame.toHtmlData( return DataFrameHtmlData(style = "", body = body, script = script) } +private const val DEFAULT_HTML_IMG_SIZE = 100 + +/** + * This method resizes a BufferedImage if necessary, according to the provided DisplayConfiguration. + * It is essential to prevent potential memory problems when serializing HTML data for display in the Kotlin Notebook plugin. + * + * @param value The input value to be checked and possibly downsized. + * @param renderConfig The DisplayConfiguration to determine if downsizing is needed. + * @return The downsized BufferedImage if value is a BufferedImage and downsizing is enabled in the DisplayConfiguration, + * otherwise returns the input value unchanged. + */ +private fun downsizeBufferedImageIfNeeded(value: Any?, renderConfig: DisplayConfiguration): Any? = + when { + value is BufferedImage && renderConfig.downsizeBufferedImage -> { + value.resizeKeepingAspectRatio(DEFAULT_HTML_IMG_SIZE) + } + + else -> value + } + /** * Renders [this] [DataFrame] as static HTML (meaning no JS is used). * CSS rendering is enabled by default but can be turned off using [includeCss] @@ -568,6 +591,7 @@ public data class DisplayConfiguration( internal val localTesting: Boolean = flagFromEnv("KOTLIN_DATAFRAME_LOCAL_TESTING"), var useDarkColorScheme: Boolean = false, var enableFallbackStaticTables: Boolean = true, + var downsizeBufferedImage: Boolean = true ) { public companion object { public val DEFAULT: DisplayConfiguration = DisplayConfiguration() diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt index e648458f9..295a40413 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt @@ -276,6 +276,7 @@ public fun AnyFrame.toJson(prettyPrint: Boolean = false, canonical: Boolean = fa * Applied for each frame column recursively * @param prettyPrint Specifies whether the output JSON should be formatted with indentation and line breaks. * @param canonical Specifies whether the output JSON should be in a canonical form. + * @param imageEncodingOptions The options for encoding images. The default is null, which indicates that the image is not encoded as Base64. * * @return The DataFrame converted to a JSON string with metadata. */ @@ -283,13 +284,39 @@ public fun AnyFrame.toJsonWithMetadata( rowLimit: Int, nestedRowLimit: Int? = null, prettyPrint: Boolean = false, - canonical: Boolean = false + canonical: Boolean = false, + imageEncodingOptions: Base64ImageEncodingOptions? = null ): String { return json { - encodeDataFrameWithMetadata(this@toJsonWithMetadata, rowLimit, nestedRowLimit) + encodeDataFrameWithMetadata(this@toJsonWithMetadata, rowLimit, nestedRowLimit, imageEncodingOptions) }.toJsonString(prettyPrint, canonical) } +internal const val DEFAULT_IMG_SIZE = 600 + +/** + * Class representing the options for encoding images. + * + * @property imageSizeLimit The maximum size to which images should be resized. Defaults to the value of DEFAULT_IMG_SIZE. + * @property options Bitwise-OR of the [GZIP_ON] and [LIMIT_SIZE_ON] constants. Defaults to [GZIP_ON] or [LIMIT_SIZE_ON]. + */ +public class Base64ImageEncodingOptions( + public val imageSizeLimit: Int = DEFAULT_IMG_SIZE, + private val options: Int = GZIP_ON or LIMIT_SIZE_ON +) { + public val isGzipOn: Boolean + get() = options and GZIP_ON == GZIP_ON + + public val isLimitSizeOn: Boolean + get() = options and LIMIT_SIZE_ON == LIMIT_SIZE_ON + + public companion object { + public const val ALL_OFF: Int = 0 + public const val GZIP_ON: Int = 1 // 2^0 + public const val LIMIT_SIZE_ON: Int = 2 // 2^1 + } +} + public fun AnyRow.toJson(prettyPrint: Boolean = false, canonical: Boolean = false): String { return json { encodeRow(df(), index()) diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/JupyterHtmlRenderer.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/JupyterHtmlRenderer.kt index 536470fa8..a5a4dc249 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/JupyterHtmlRenderer.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/JupyterHtmlRenderer.kt @@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.jupyter import com.beust.klaxon.json import org.jetbrains.kotlinx.dataframe.api.take import org.jetbrains.kotlinx.dataframe.impl.io.encodeFrame +import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions import org.jetbrains.kotlinx.dataframe.io.DataFrameHtmlData import org.jetbrains.kotlinx.dataframe.io.DisplayConfiguration import org.jetbrains.kotlinx.dataframe.io.toHTML @@ -23,6 +24,7 @@ import org.jetbrains.kotlinx.jupyter.api.renderHtmlAsIFrameIfNeeded /** Starting from this version, dataframe integration will respond with additional data for rendering in Kotlin Notebooks plugin. */ private const val MIN_KERNEL_VERSION_FOR_NEW_TABLES_UI = "0.11.0.311" private const val MIN_IDE_VERSION_SUPPORT_JSON_WITH_METADATA = 241 +private const val MIN_IDE_VERSION_SUPPORT_IMAGE_VIEWER = 242 internal class JupyterHtmlRenderer( val display: DisplayConfiguration, @@ -63,8 +65,8 @@ internal inline fun JupyterHtmlRenderer.render( if (notebook.kernelVersion >= KotlinKernelVersion.from(MIN_KERNEL_VERSION_FOR_NEW_TABLES_UI)!!) { val ideBuildNumber = KotlinNotebookPluginUtils.getKotlinNotebookIDEBuildNumber() - val jsonEncodedDf = - if (ideBuildNumber == null || ideBuildNumber.majorVersion < MIN_IDE_VERSION_SUPPORT_JSON_WITH_METADATA) { + val jsonEncodedDf = when { + !ideBuildNumber.supportsDynamicNestedTables() -> { json { obj( "nrow" to df.size.nrow, @@ -73,15 +75,32 @@ internal inline fun JupyterHtmlRenderer.render( "kotlin_dataframe" to encodeFrame(df.take(limit)), ) }.toJsonString() - } else { - df.toJsonWithMetadata(limit, reifiedDisplayConfiguration.rowsLimit) } + + else -> { + val imageEncodingOptions = + if (ideBuildNumber.supportsImageViewer()) Base64ImageEncodingOptions() else null + + df.toJsonWithMetadata( + limit, + reifiedDisplayConfiguration.rowsLimit, + imageEncodingOptions = imageEncodingOptions + ) + } + } + notebook.renderAsIFrameAsNeeded(html, staticHtml, jsonEncodedDf) } else { notebook.renderHtmlAsIFrameIfNeeded(html) } } +private fun KotlinNotebookPluginUtils.IdeBuildNumber?.supportsDynamicNestedTables() = + this != null && majorVersion >= MIN_IDE_VERSION_SUPPORT_JSON_WITH_METADATA + +private fun KotlinNotebookPluginUtils.IdeBuildNumber?.supportsImageViewer() = + this != null && majorVersion >= MIN_IDE_VERSION_SUPPORT_IMAGE_VIEWER + internal fun Notebook.renderAsIFrameAsNeeded( data: HtmlData, staticData: HtmlData, diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/Utils.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/Utils.kt index c071955e0..dae5f687b 100644 --- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/Utils.kt +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/Utils.kt @@ -1,5 +1,7 @@ package org.jetbrains.kotlinx.dataframe +import com.beust.klaxon.JsonObject +import com.beust.klaxon.Parser import org.jetbrains.kotlinx.dataframe.api.print import org.jetbrains.kotlinx.dataframe.api.schema import org.jetbrains.kotlinx.dataframe.io.renderToString @@ -24,3 +26,8 @@ fun > T.alsoDebug(println: String? = null, rowsLimit: Int = 20) print(borders = true, title = true, columnTypes = true, valueLimit = -1, rowsLimit = rowsLimit) schema().print() } + +fun parseJsonStr(jsonStr: String): JsonObject { + val parser = Parser.default() + return parser.parse(StringBuilder(jsonStr)) as JsonObject +} diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/ImageSerializationTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/ImageSerializationTests.kt new file mode 100644 index 000000000..cae6e759c --- /dev/null +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/ImageSerializationTests.kt @@ -0,0 +1,174 @@ +package org.jetbrains.kotlinx.dataframe.io + +import com.beust.klaxon.JsonArray +import com.beust.klaxon.JsonObject +import io.kotest.matchers.shouldBe +import io.kotest.matchers.string.shouldContain +import org.jetbrains.kotlinx.dataframe.api.dataFrameOf +import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.KOTLIN_DATAFRAME +import org.jetbrains.kotlinx.dataframe.impl.io.resizeKeepingAspectRatio +import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions.Companion.ALL_OFF +import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions.Companion.GZIP_ON +import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions.Companion.LIMIT_SIZE_ON +import org.jetbrains.kotlinx.dataframe.parseJsonStr +import org.jetbrains.kotlinx.dataframe.testResource +import org.junit.Test +import org.junit.runner.RunWith +import org.junit.runners.Parameterized +import java.awt.image.BufferedImage +import java.io.ByteArrayInputStream +import java.io.ByteArrayOutputStream +import java.io.File +import java.util.* +import java.util.zip.GZIPInputStream +import javax.imageio.ImageIO + +@RunWith(Parameterized::class) +class ImageSerializationTests(private val encodingOptions: Base64ImageEncodingOptions?) { + @Test + fun `serialize images as base64`() { + val images = readImagesFromResources() + val json = encodeImagesAsJson(images, encodingOptions) + + if (encodingOptions == DISABLED) { + checkImagesEncodedAsToString(json, images.size) + return + } + + val decodedImages = decodeImagesFromJson(json, images.size, encodingOptions!!) + + for ((decodedImage, original) in decodedImages.zip(images)) { + val expectedImage = resizeIfNeeded(original, encodingOptions) + isImagesIdentical(decodedImage, expectedImage, 2) shouldBe true + } + } + + private fun readImagesFromResources(): List { + val dir = File(testResource("imgs").path) + + return dir.listFiles()?.map { file -> + try { + ImageIO.read(file) + } catch (ex: Exception) { + throw IllegalArgumentException("Error reading ${file.name}: ${ex.message}") + } + } ?: emptyList() + } + + private fun encodeImagesAsJson( + images: List, + encodingOptions: Base64ImageEncodingOptions? + ): JsonObject { + val df = dataFrameOf(listOf("imgs"), images) + val jsonStr = df.toJsonWithMetadata(20, nestedRowLimit = 20, imageEncodingOptions = encodingOptions) + + return parseJsonStr(jsonStr) + } + + private fun checkImagesEncodedAsToString(json: JsonObject, numImgs: Int) { + for (i in 0..)[i] as JsonObject + val img = row["imgs"] as String + + img shouldContain "BufferedImage" + } + } + + private fun decodeImagesFromJson( + json: JsonObject, + imgsNum: Int, + encodingOptions: Base64ImageEncodingOptions + ): List { + val result = mutableListOf() + for (i in 0..)[i] as JsonObject + val imgString = row["imgs"] as String + + val bytes = decodeBase64Image(imgString, encodingOptions) + val decodedImage = createImageFromBytes(bytes) + + result.add(decodedImage) + } + + return result + } + + private fun decodeBase64Image(imgString: String, encodingOptions: Base64ImageEncodingOptions): ByteArray = + when { + encodingOptions.isGzipOn -> decompressGzip(Base64.getDecoder().decode(imgString)) + else -> Base64.getDecoder().decode(imgString) + } + + private fun decompressGzip(input: ByteArray): ByteArray { + return ByteArrayOutputStream().use { byteArrayOutputStream -> + GZIPInputStream(input.inputStream()).use { inputStream -> + inputStream.copyTo(byteArrayOutputStream) + } + byteArrayOutputStream.toByteArray() + } + } + + private fun resizeIfNeeded(image: BufferedImage, encodingOptions: Base64ImageEncodingOptions): BufferedImage = + when { + !encodingOptions.isLimitSizeOn -> image + else -> image.resizeKeepingAspectRatio(encodingOptions.imageSizeLimit) + } + + private fun createImageFromBytes(bytes: ByteArray): BufferedImage { + val bais = ByteArrayInputStream(bytes) + return ImageIO.read(bais) + } + + private fun isImagesIdentical(img1: BufferedImage, img2: BufferedImage, allowedDelta: Int): Boolean { + // First check dimensions + if (img1.width != img2.width || img1.height != img2.height) { + return false + } + + // Then check each pixel + for (y in 0 until img1.height) { + for (x in 0 until img1.width) { + val rgb1 = img1.getRGB(x, y) + val rgb2 = img2.getRGB(x, y) + + val r1 = (rgb1 shr 16) and 0xFF + val g1 = (rgb1 shr 8) and 0xFF + val b1 = rgb1 and 0xFF + + val r2 = (rgb2 shr 16) and 0xFF + val g2 = (rgb2 shr 8) and 0xFF + val b2 = rgb2 and 0xFF + + val diff = kotlin.math.abs(r1 - r2) + kotlin.math.abs(g1 - g2) + kotlin.math.abs(b1 - b2) + + // If the difference in color components exceed our allowance return false + if (diff > allowedDelta) { + return false + } + } + } + + // If no exceeding difference was found, the images are identical within our allowedDelta + return true + } + + companion object { + private val DEFAULT = Base64ImageEncodingOptions() + private val GZIP_ON_RESIZE_OFF = Base64ImageEncodingOptions(options = GZIP_ON) + private val GZIP_OFF_RESIZE_OFF = Base64ImageEncodingOptions(options = ALL_OFF) + private val GZIP_ON_RESIZE_TO_700 = Base64ImageEncodingOptions(imageSizeLimit = 700, options = GZIP_ON or LIMIT_SIZE_ON) + private val DISABLED = null + + @JvmStatic + @Parameterized.Parameters + fun imageEncodingOptionsToTest(): Collection { + return listOf( + DEFAULT, + GZIP_ON_RESIZE_OFF, + GZIP_OFF_RESIZE_OFF, + GZIP_ON_RESIZE_TO_700, + null, + ) + } + } +} diff --git a/core/src/test/resources/imgs/img1.jpg b/core/src/test/resources/imgs/img1.jpg new file mode 100644 index 000000000..40d5662e4 Binary files /dev/null and b/core/src/test/resources/imgs/img1.jpg differ diff --git a/core/src/test/resources/imgs/img2.jpg b/core/src/test/resources/imgs/img2.jpg new file mode 100644 index 000000000..e8e34329c Binary files /dev/null and b/core/src/test/resources/imgs/img2.jpg differ diff --git a/core/src/test/resources/imgs/img3.jpg b/core/src/test/resources/imgs/img3.jpg new file mode 100644 index 000000000..33864500a Binary files /dev/null and b/core/src/test/resources/imgs/img3.jpg differ