PngIOMetadata.java

/* ========================================================================
 * PlantUML : a free UML diagram generator
 * ========================================================================
 *
 * (C) Copyright 2009-2024, Arnaud Roques
 *
 * Project Info:  https://plantuml.com
 * 
 * If you like this project or if you find it useful, you can support us at:
 * 
 * https://plantuml.com/patreon (only 1$ per month!)
 * https://plantuml.com/paypal
 * 
 * This file is part of PlantUML.
 *
 * PlantUML is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * PlantUML distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public
 * License for more details.
 *
 * You should have received a copy of the GNU General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301,
 * USA.
 *
 *
 * Original Author:  Arnaud Roques
 * 
 *
 */
package net.sourceforge.plantuml.png;

import java.awt.image.BufferedImage;
import java.awt.image.RenderedImage;
import java.io.IOException;
import java.io.OutputStream;

import javax.imageio.IIOImage;
import javax.imageio.ImageTypeSpecifier;
import javax.imageio.ImageWriteParam;
import javax.imageio.ImageWriter;
import javax.imageio.metadata.IIOInvalidTreeException;
import javax.imageio.metadata.IIOMetadata;
import javax.imageio.metadata.IIOMetadataNode;
import javax.imageio.stream.ImageOutputStream;

import net.sourceforge.plantuml.directdot.CounterOutputStream;
import net.sourceforge.plantuml.security.SImageIO;
import net.sourceforge.plantuml.utils.Log;

public class PngIOMetadata {
	// ::remove file when __CORE__ or __TEAVM__

	private static final String copyleft = "Generated by https://plantuml.com";

	public static void writeWithMetadata(RenderedImage image, OutputStream os, String metadata, int dpi,
			String debugData, int level) throws IOException {

		final int w = image.getWidth();
		final int h = image.getHeight();
		final String imageType = image.getClass().getSimpleName();
		final String colorModel = image instanceof BufferedImage ? getImageTypeName(((BufferedImage) image).getType())
				: "unknown";

		Log.info(() -> "PngIOMetadata: starting, image " + w + "x" + h + " (" + imageType + ", " + colorModel
				+ "), compression level=" + level + ", memory: " + PngIO.getUsedMemoryMB() + " MB");

		final long startTotal = System.currentTimeMillis();
		final long memBefore = PngIO.getUsedMemoryMB();

		final ImageWriter writer = javax.imageio.ImageIO.getImageWritersByFormatName("png").next();
		try {
			final ImageWriteParam writeParam = writer.getDefaultWriteParam();

			try {
				writeParam.setCompressionMode(ImageWriteParam.MODE_EXPLICIT);
				writeParam.setCompressionQuality(levelToQuality(level));
			} catch (Throwable t) {
				Log.debug(() -> "Warning: cannot set compression mode/quality");
			}

			final ImageTypeSpecifier typeSpecifier = ImageTypeSpecifier
					.createFromBufferedImageType(BufferedImage.TYPE_INT_RGB);

			final IIOMetadata meta = writer.getDefaultImageMetadata(typeSpecifier, writeParam);

			if (dpi != 96)
				addDpi(meta, dpi);

			if (debugData != null)
				addText(meta, "debug", debugData);

			addText(meta, "copyleft", copyleft);
			if (metadata != null)
				addiText(meta, "plantuml", metadata);

			Log.debug(() -> "PngIOMetadata pngMetadata=" + meta);

			// Render the PNG to file
			final IIOImage iioImage = new IIOImage(image, null, meta);
			Log.debug(() -> "PngIOMetadata iioImage=" + iioImage);

			final long startWrite = System.currentTimeMillis();
			final CounterOutputStream cos = new CounterOutputStream(os);
			try (final ImageOutputStream ios = SImageIO.createImageOutputStream(cos)) {
				writer.setOutput(ios);
				writer.write(null, iioImage, writeParam);
				ios.flush();

				final long writeDuration = System.currentTimeMillis() - startWrite;
				final long totalDuration = System.currentTimeMillis() - startTotal;
				final long memAfter = PngIO.getUsedMemoryMB();
				final int byteCount = cos.getLength();
				Log.info(() -> "PngIOMetadata: PNG written, " + byteCount + " bytes (" + formatSize(byteCount)
						+ "), write: " + writeDuration + " ms, total: " + totalDuration + " ms, memory: " + memBefore
						+ " -> " + memAfter + " MB (delta: " + (memAfter - memBefore) + " MB)");
			}
		} finally {
			writer.dispose();
		}
	}

	private static String formatSize(long bytes) {
		if (bytes < 1024)
			return bytes + " B";
		if (bytes < 1024 * 1024)
			return String.format("%.1f KB", bytes / 1024.0);
		return String.format("%.2f MB", bytes / (1024.0 * 1024.0));
	}

	private static String getImageTypeName(int type) {
		switch (type) {
		case BufferedImage.TYPE_INT_RGB:
			return "INT_RGB";
		case BufferedImage.TYPE_INT_ARGB:
			return "INT_ARGB";
		case BufferedImage.TYPE_INT_ARGB_PRE:
			return "INT_ARGB_PRE";
		case BufferedImage.TYPE_INT_BGR:
			return "INT_BGR";
		case BufferedImage.TYPE_3BYTE_BGR:
			return "3BYTE_BGR";
		case BufferedImage.TYPE_4BYTE_ABGR:
			return "4BYTE_ABGR";
		case BufferedImage.TYPE_4BYTE_ABGR_PRE:
			return "4BYTE_ABGR_PRE";
		case BufferedImage.TYPE_BYTE_GRAY:
			return "BYTE_GRAY";
		case BufferedImage.TYPE_BYTE_BINARY:
			return "BYTE_BINARY";
		case BufferedImage.TYPE_BYTE_INDEXED:
			return "BYTE_INDEXED";
		case BufferedImage.TYPE_USHORT_GRAY:
			return "USHORT_GRAY";
		case BufferedImage.TYPE_USHORT_565_RGB:
			return "USHORT_565_RGB";
		case BufferedImage.TYPE_USHORT_555_RGB:
			return "USHORT_555_RGB";
		case BufferedImage.TYPE_CUSTOM:
			return "CUSTOM";
		default:
			return "type=" + type;
		}
	}

	private static float levelToQuality(int level) {
		final int L = Math.max(1, Math.min(9, level));
		return 1.0f - (L - 1) / 8.0f;
	}

	private static void addDpi(IIOMetadata meta, double dpi) throws IIOInvalidTreeException {
		final IIOMetadataNode dimension = new IIOMetadataNode("Dimension");

		final IIOMetadataNode horizontalPixelSize = new IIOMetadataNode("HorizontalPixelSize");
		final double value = dpi / 0.0254 / 1000;
		horizontalPixelSize.setAttribute("value", Double.toString(value));
		dimension.appendChild(horizontalPixelSize);

		final IIOMetadataNode verticalPixelSize = new IIOMetadataNode("VerticalPixelSize");
		verticalPixelSize.setAttribute("value", Double.toString(value));
		dimension.appendChild(verticalPixelSize);

		final IIOMetadataNode root = new IIOMetadataNode("javax_imageio_1.0");
		root.appendChild(dimension);

		meta.mergeTree("javax_imageio_1.0", root);

	}

	private static void addiText(IIOMetadata meta, String key, String value) throws IIOInvalidTreeException {
		final IIOMetadataNode text = new IIOMetadataNode("iTXt");
		final IIOMetadataNode entry = new IIOMetadataNode("iTXtEntry");
		entry.setAttribute("keyword", key);
		entry.setAttribute("compressionFlag", "true");
		entry.setAttribute("compressionMethod", "0");
		entry.setAttribute("languageTag", "");
		entry.setAttribute("translatedKeyword", "");
		entry.setAttribute("text", value);

		text.appendChild(entry);
		final IIOMetadataNode root = new IIOMetadataNode("javax_imageio_png_1.0");
		root.appendChild(text);

		meta.mergeTree("javax_imageio_png_1.0", root);

	}

	private static void addText(IIOMetadata meta, String key, String value) throws IIOInvalidTreeException {
		final IIOMetadataNode text = new IIOMetadataNode("tEXt");
		final IIOMetadataNode entry = new IIOMetadataNode("tEXtEntry");
		entry.setAttribute("keyword", key);
		entry.setAttribute("value", value);

		text.appendChild(entry);
		final IIOMetadataNode root = new IIOMetadataNode("javax_imageio_png_1.0");
		root.appendChild(text);

		meta.mergeTree("javax_imageio_png_1.0", root);
	}

}