Displays all clusters

gs-cv/src/main/java/org/genericsystem/cv/Kmeans.java

@@ -4,9 +4,7 @@
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 
 import javax.imageio.ImageIO;
 import javax.swing.ImageIcon;
@@ -19,6 +17,7 @@
 import org.opencv.core.MatOfByte;
 import org.opencv.core.TermCriteria;
 import org.opencv.imgcodecs.Imgcodecs;
+import org.opencv.utils.Converters;
 import org.opencv.videoio.VideoCapture;
 
 public class Kmeans {
@@ -28,6 +27,9 @@ public class Kmeans {
 	}
 	private static VideoCapture camera = new VideoCapture(0);
 
+	private final static int MAX_ITER = 5;
+	private final static int CLUSTERS = 15;
+
 	public static void main(String[] args) {
 		JFrame jframe = new JFrame("Kmeans");
 		jframe.setResizable(false);
@@ -37,60 +39,38 @@ public static void main(String[] args) {
 
 		Mat img = new Mat();
 		camera.read(img);
-		// Mat img = Mat.zeros(200, 200, CvType.CV_8UC3);
-		// Imgproc.rectangle(img, new Point(0, 0), new Point(100, 200), new Scalar(0, 255, 0), -1);
-		// Imgproc.rectangle(img, new Point(100, 0), new Point(200, 200), new Scalar(0, 0, 255), -1);
-
-		Mat clusters = cluster(img, 10).get(2);
+		jframe.setSize(img.width(), img.height());
+		jframe.setVisible(true);
+		for (;;) {
+			camera.read(img);
+			Mat clusters = colorMapKMeans(img, CLUSTERS, MAX_ITER);
 
-		ImageIcon image = new ImageIcon(mat2bufferedImage(clusters));
+			ImageIcon image = new ImageIcon(mat2bufferedImage(clusters));
 
-		jframe.setSize(clusters.width(), clusters.height());
-		jframe.setVisible(true);
-		vidpanel.setIcon(image);
-		vidpanel.repaint();
+			vidpanel.setIcon(image);
+			vidpanel.repaint();
+		}
 
 	}
 
-	public static List<Mat> cluster(Mat cutout, int k) {
-		Mat samples = cutout.reshape(1, cutout.cols() * cutout.rows());
-		Mat samples32f = new Mat();
-		samples.convertTo(samples32f, CvType.CV_32F, 1.0 / 255.0);
+	public static Mat colorMapKMeans(Mat img, int K, int maxIterations) {
 
-		Mat labels = new Mat();
-		TermCriteria criteria = new TermCriteria(TermCriteria.COUNT, 100, 1);
-		Mat centers = new Mat();
-		Core.kmeans(samples32f, k, labels, criteria, 1, Core.KMEANS_PP_CENTERS, centers);
-		return showClusters(cutout, labels, centers);
-	}
+		Mat m = img.reshape(1, img.rows() * img.cols());
+		m.convertTo(m, CvType.CV_32F);
 
-	private static List<Mat> showClusters(Mat cutout, Mat labels, Mat centers) {
-		centers.convertTo(centers, CvType.CV_8UC1, 255.0);
-		centers.reshape(3);
+		Mat bestLabels = new Mat(m.rows(), 1, CvType.CV_8U);
+		Mat centroids = new Mat(K, 1, CvType.CV_32F);
+		Core.kmeans(m, K, bestLabels, new TermCriteria(TermCriteria.COUNT | TermCriteria.EPS, maxIterations, 1E-5), 1, Core.KMEANS_RANDOM_CENTERS, centroids);
+		List<Integer> idx = new ArrayList<>(m.rows());
+		Converters.Mat_to_vector_int(bestLabels, idx);
 
-		List<Mat> clusters = new ArrayList<Mat>();
-		for (int i = 0; i < centers.rows(); i++) {
-			clusters.add(Mat.zeros(cutout.size(), cutout.type()));
+		Mat imgMapped = new Mat(m.size(), m.type());
+		for (int i = 0; i < idx.size(); i++) {
+			Mat row = imgMapped.row(i);
+			centroids.row(idx.get(i)).copyTo(row);
 		}
 
-		Map<Integer, Integer> counts = new HashMap<Integer, Integer>();
-		for (int i = 0; i < centers.rows(); i++)
-			counts.put(i, 0);
-
-		int rows = 0;
-		for (int y = 0; y < cutout.rows(); y++) {
-			for (int x = 0; x < cutout.cols(); x++) {
-				int label = (int) labels.get(rows, 0)[0];
-				int r = (int) centers.get(label, 2)[0];
-				int g = (int) centers.get(label, 1)[0];
-				int b = (int) centers.get(label, 0)[0];
-				counts.put(label, counts.get(label) + 1);
-				clusters.get(label).put(y, x, b, g, r);
-				rows++;
-			}
-		}
-		System.out.println(counts);
-		return clusters;
+		return imgMapped.reshape(3, img.rows());
 	}
 
 	public static BufferedImage mat2bufferedImage(Mat image) {