net.sourceforge.plantuml.quantization.KMeansQuantizer Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of plantuml-mit Show documentation

PlantUML is a component that allows to quickly write diagrams from text.

There is a newer version: 1.2024.8

// THIS FILE HAS BEEN GENERATED BY A PREPROCESSOR.
package net.sourceforge.plantuml.quantization;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Uses k-means clustering for color quantization. This tends to yield good
 * results, but convergence can be slow. It is not recommended for large images.
 */
public final class KMeansQuantizer implements ColorQuantizer {
	public static final KMeansQuantizer INSTANCE = new KMeansQuantizer();

	private KMeansQuantizer() {
	}

	@Override
	public Set quantize(Multiset originalColors, int maxColorCount) {
		Map> clustersByCentroid = new LinkedHashMap<>();
		Set centroidsToRecompute = getInitialCentroids(originalColors, maxColorCount);
		for (QColor centroid : centroidsToRecompute)
			clustersByCentroid.put(centroid, new HashMultiset());

		for (QColor color : originalColors.getDistinctElements()) {
			final int count = originalColors.count(color);
			clustersByCentroid.get(color.getNearestColor(centroidsToRecompute)).add(color, count);
		}

		while (centroidsToRecompute.isEmpty() == false) {
			recomputeCentroids(clustersByCentroid, centroidsToRecompute);
			centroidsToRecompute.clear();

			Set allCentroids = clustersByCentroid.keySet();
			for (QColor centroid : clustersByCentroid.keySet()) {
				Multiset cluster = clustersByCentroid.get(centroid);
				for (QColor color : new ArrayList<>(cluster.getDistinctElements())) {
					QColor newCentroid = color.getNearestColor(allCentroids);
					if (newCentroid != centroid) {
						final int count = cluster.count(color);
						final Multiset newCluster = clustersByCentroid.get(newCentroid);

						cluster.remove(color, count);
						newCluster.add(color, count);

						centroidsToRecompute.add(centroid);
						centroidsToRecompute.add(newCentroid);
					}
				}
			}
		}

		return clustersByCentroid.keySet();
	}

	private static void recomputeCentroids(Map> clustersByCentroid,
			Set centroidsToRecompute) {
		for (QColor oldCentroid : centroidsToRecompute) {
			final Multiset cluster = clustersByCentroid.get(oldCentroid);
			final QColor newCentroid = QColor.getCentroid(cluster);
			clustersByCentroid.remove(oldCentroid);
			clustersByCentroid.put(newCentroid, cluster);
		}
	}

	private static Set getInitialCentroids(Multiset originalColors, int maxColorCount) {
		// We use the Forgy initialization method: choose random colors as initial
		// cluster centroids.
		final List colorList = new ArrayList<>(originalColors.getDistinctElements());
		Collections.shuffle(colorList);
		return new HashSet<>(colorList.subList(0, maxColorCount));
	}
}