Image Segmentation on Apache Spark
By: Geraldo Souza Junior
Tags: Machine Learning, Scala, Image Segmentation, K Means Clustering, Apache Spark

Image Segmentation on Apache Spark

Source: Unsplash

Computer Vision is one of the most exciting branches of data science. There is many possible applications for applying Machine Learning algorithms and techniques and Image Segmentation is one of the first steps into it.

Image Segmentation is one of the main developments for processing on Computer Vision. Many Computer Vision tasks requires the segmentation of an image, to understand each part and easier comprehension as a whole. Each segment contains a set of pixels, which may represent something. The segmentation is the basis for two important applications of image analysis:

Object Detection: identify objects in the image, cars, people or animal, etc. The goal here is to just identify that there is an object in that set of pixels, but not to classify properly which object was it.

identify objects in the image, cars, people or animal, etc. The goal here is to just identify that there is an object in that set of pixels, but not to classify properly which object was it. Classification: after the object detection, the objective is to classify each detected object into a class. The image below illustrates this application, with each object detected and classified.

Image Source: Towards a Meaningful 3D Map Using a 3D Lidar and a Camera. https://www.researchgate.net/publication/326875064_Towards_a_Meaningful_3D_Map_Using_a_3D_Lidar_and_a_Camera

Segmentation Techniques

There is many ways to achieve the segmentation of an image, through Deep Learning, region based approaches (which will try to detect the borders of an object), through the application of clustering algorithms and so on. Each technique as its advantages and disadvantages. Deep Learning techniques has proven the best results so far, but the implementation tends to be complex and requires a large number of data to train the model. Of course, it is possible to use pre trained models available, but for use on a specific application, it will require training, which can be time and processing consuming. Another possible way, is to apply clustering algorithms. The goal is to identify groups in the data, assigning desired information from each pixel (x and y positions, color information), to classify them in one of the groups, based on feature similarity. In contrast to Deep Learning, it is less efficient, because rather than analyzing pre defined groups, clustering works iteratively to organically form groups, which will require more human interference. However, since clustering is unsupervised, it does not require too much data too train, which can reduce development time and save processing time compared to Deep Learning counterparts.

K-Means Clustering

Before organize the data into groups, we need to understand the concept of distance between two data points, in analytic geometry. There is many ways to compute distance, but i will focus here in one of the most commonly used methods: Euclidean distance. According to wikipedia ,the Euclidean distance between two points p and q is the length of the line segment connecting them. If p = (p1, p2,…, pn) and q = (q1, q2,…, qn) are two points in Euclidean n-space, the Pythagorean Theorem can be used to calculate the distance between two points:

Let’s use a practical example. Given two points p and q , where p=(5,3) and q = (4,2). Using the image below as illustration, x1 = 5, y1 =3, x2 = 4, y2 = 2. Applying the formula, we have d(p,q) = (4–5)2+ (2–3)2 = 2.

Now that we have an understanding of distance, we can proceed to K-Means computation. First it is defined first the final number of clusters K. The number of cluster should be defined considering the application context of each case. After that, it is necessary to define the centroids for the clusters. A centroid is defined as the average of all data points in a given cluster. So, if a data point P is closer to a centroid K1 instead of K2, we can say that this point belongs to the group K1 . There is a lot of ways to initialize the centroids. Here it is some of then:

RP — randomly selected points. k distinct cases of the data are randomly selected to be the initial centres.

— randomly selected points. k distinct cases of the data are randomly selected to be the initial centres. KMPP — random farthest points, or k-means++. The first centre is selected as a random case from the dataset. The 2nd centre is selected also randomly, but the probability of selection of a case is proportional to the distance (square euclidean) of it to that (1st) centre. The 3rd centre is selected also randomly with the probability of selection proportional to the distance of a case to the nearest of those two centres, — and so on.

— random farthest points, or k-means++. The first centre is selected as a random case from the dataset. The 2nd centre is selected also randomly, but the probability of selection of a case is proportional to the distance (square euclidean) of it to that (1st) centre. The 3rd centre is selected also randomly with the probability of selection proportional to the distance of a case to the nearest of those two centres, — and so on. SIMFP — farthest points (simple selection). The first centre is selected as a random case from the dataset. The 2nd centre is selected as the case maximally distant from that centre. The 3rd centre is selected as the case maximally distant from those two (from the nearest of the two), — and so on.

— farthest points (simple selection). The first centre is selected as a random case from the dataset. The 2nd centre is selected as the case maximally distant from that centre. The 3rd centre is selected as the case maximally distant from those two (from the nearest of the two), — and so on. RUNFP — farthest points (running selection). First k cases are taken as centres and then during the run through the rest of the cases of the dataset there progressively replacements among the centres are done; the aim of the replacements is to obtain in the end k points most distant from each other in the variable space. These points (cases) occupying peripheral positions in the data cloud are the produced initial centres.

— farthest points (running selection). First k cases are taken as centres and then during the run through the rest of the cases of the dataset there progressively replacements among the centres are done; the aim of the replacements is to obtain in the end k points most distant from each other in the variable space. These points (cases) occupying peripheral positions in the data cloud are the produced initial centres. GREP — group representative points. The method idea — to collect as centres k most representative, “deputy” cases. The 1st centre is taken as the case closest to the general data cenroid. Then the rest of the centres are selected from the data points in such a way that each point is considered as to whether it is closer (and how much, in terms of squared euclidean distance) to a set of points than each one of the latter is to any of the already existing centres. I.e. each point is examined as a candidate to represent some group of points not yet well enough represented by the centres already collected. Point most representative in this respect is selected as the next centre.

So, now that we have now that we have consolidated the distance and the initialization of centroids, let’s see how the K-Means algorithms works:

If you still have any doubts in any step, i fully recommend this step by step calculation, wrote by Sunaina, on datasciencecentral. Also, this video shows a visualization of each step.

Implementing Image Segmentation with K-Mean on Spark

Ok, time to start solving the image segmentation problem with the k-means clustering algorithm on apache spark with scala. Wait, but why scala? Python is currently the most preferred language among the data scientists not just it is easy to learn and implement but also for its extensive libraries and frameworks. In data science and machine learning projects, it includes a broad range of useful libraries SciPy, NumPy, Matplolib, Pandas, among others while for more complex projects in deep learning, Python offers libraries such as Keras, Pytorch, and TensorFlow. With all these advantages, again, why scala? Well, Spark is written in Scala as it can be quite fast because it’s statically typed and it compiles in a known way to the JVM. I came from the Java world, and one of the advantages is to use the whole legacy from Java, such as libraries and frameworks.

I will not dig into the differences between languages because this is not the point of this post (there is hundreds of pos and cons available, google is your friend), but like said before, a fact is that the Python community is huge and provides many libs and frameworks to use together with Spark, a lot of tutorials available. Because of this, my main motivation with this post, is to contribute with the scala community.

Loading Images into Spark Data Frames

Since, the version 2.4, Spark has a new data source built in, to read compressed formats (jpg, png, etc). The image will be read with the ImageIO Java Library, and it has a special data frame schema . Reading an image file (or several, just pointing to a direcory that contains only images) can be done by:

The schema contains a StructType Column “Image”, that contains all the information about the read data. Inside the StructType, you can find the following columns:

origin: StringType (represents the file path of the image)

(represents the file path of the image) height: IntegerType (height of the image)

(height of the image) width: IntegerType (width of the image)

(width of the image) nChannels: IntegerType (number of image channels)

(number of image channels) mode: IntegerType (OpenCV-compatible type)

(OpenCV-compatible type) data: BinaryType (Image bytes in OpenCV-compatible order: row-wise BGR in most cases)

The important data for us, is the array of bytes that spark generates for us, representing the raw image data. The array contains all the features that we may use for clustering.

Digital images are composed by pixels, and each pixel is composed by a combinations of colors, represented by code. One grayscale image for example, has only one channel of color. For computer displays, the most used standard, is the RGB. An RGB image has three channels, each one meaning the intensivity brightness of the respective colours: red, green and blue. So, basically, an RGB image is a combination of three images (one for each channel), each one containing information about the respective colour.

Have you ever tried too look closer into old CRT tv’s screens? If so, you probably remember that you could see the three color channels. The ideia is the same here, as illustrates the image .

How Spark decode the image

To build our feature map, we need to understand how spark read and decode the image. Diving into spark code, we have the ImageSchema.scala object, more specifically, the decode method.

This method basically checks the number of channels that the image has, and decodes based on this number. Let’s take an example: consider an 640 x 480 image, which will have 307200 pixels. If the image is grayscaled (only 1 channel), the final array size will be 307200. However, if it is an RGB image, the final array will have 307200 * 3 = 921600. The Channel order is defined as BGR (and also there is space for the alpha channel, which define transparency). With the same 640 x 480 RGB image example, the first three positions of the array will contain color information(Blue, Green and Red) from the pixel at height = 0 and width = 0. The next three position will contain information from the pixel at at height = 0 and width = 1 and so on.

Creating the features Data Frame

Now that we know how spark stores image information, we can read the array of data to build our features data frame, that will be used on the K-Means algorithm. For each pixel, we should have the w position, h position and each one of the color channels.

To start this task, we will use two images, as provided below, one landscape and a person(i blurred the woman’s face, just to prevent any problems).

Full size image can be found on: https://ibb.co/d4xWxvw

Full sizeimage can be found on: https://ibb.co/d4xWxvw

The logic to extract the information, is to read the array following the logic on the decode method, and them transposing it into data frame columns. Let’s follow step by step:

We start by reading the image and extracting the array of bytes and creating a new RDD for it, applying the flatMap transformation, returning each array position as an entry on the RDD. We could work directly on the image array, but we would lost one of the advatanges of using spark: distributed data structures. By working always with RDD structures, we make sure that our code will always be distributed.

The next step is to extract informations about channels, height, width and start the accumulators. Accumulators are variables to aggregate values of all spark executors. As we are using RDD’s, we cannot use regular aggregation variables, because when Spark ships this code to every executor the variables become local to that executor and its updated value is not relayed back to the driver. So, if we want to our counters work correctly (for h and w positions and the colours offset), we need to use accumulators.

The final processing begins with the zipWithIndex transformation, which will return a new RDD of pairs, a tuple with the element and an index, that will be useful for the bytes offset. The zipWithIndex will be useful to assign a index to each entry on the RDD, that will be used for the offset. In the end, the RDD will be converted into a Data Frame with three columns, representing the color information and the x and y positions. It is important to notice, that each pixel will have the number of channels described on the schema. So, it will be important to collapse all the color information of a certain pixel into only one entry of the data set. For that, it is used an aggregation on the Data Frame, grouping by the X and Y positions of the pixel. Next, it is applied the org.apache.spark.sql.functions.collect_list function, which will take all the aggregated values from the color column and return a list containing it. Finally, it is returned a DataFrame with the X and Y positions, blue, green and red color informations.

The implementation above considers only images with three channels (RGB), since it is the most common found on images. Images with 4 channels (RGBa) uses the fourth channel for transparency information and would require little adjustment in the offset position of the pixels.

The Machine Learning Pipeline on Spark

In general, most of the implementations of machine learning models can be designed as a ordered sequence of some algorithms, like the following:

Feature extraction, transformation and selection.

Train a predictive model based on the these vectors and label.

Make predictions using the generated model.

Evaluate model (performance and accuracy).

Spark MLib provides two top level abstractions to facilitate the development of this pipeline: transformers and estimators. A transformer implements a method transform() which will convert one DataFrame into another, generally appending one or more new column. For example, a transformer will take all the columns features of each entries on the Data Frame and map it into a new column (feature vectors). The estimator will be responsible for applying the learning algorithm that fits or trains on data. It implements the method fit(), that accepts a DataFrame and produces a Model, which is a transformer. As an example, a learning algorithm such as DecisionTree is an estimator, and calling fit() trains a DecisionTreeClassificationModel, which is a Model and hence a transformer.

With our DataFrame with feature columns (b,g,r,w,h) ready, we can start our pipeline by applying the transformer VectorAssembler. This transformer combines a given list of columns into a single vector column. We use the method setInputCols to pass an array, pointing the columns that we want to combine. The method setOutputCol is used to point the column that will contain the combined vector.

Initially, only the color informations will be used as features of our model. We can test and discuss the inclusion of x and y positions later. This will be the resulting DataFrame:

Now, we can create the Kmeans estimator and train our model. The variable K will define the number of clusters that will be used. The setSeed method, is for the random initialization of the clusters. If you do not provide the seed, spark will generate one internally, which will compromise tests comparisons, because the performance will be different for each run, since the clusters initial points will be different for each run. The method setFeaturesCol receives the column that contains the vector of features.

The model is created by calling the fit() method and passing the DataFrame that was transformed by the Vector Assembler. The model will give the results by invoking the transform() method, returning a DataFrame containing a new column, prediction.

Finally, the resulting DataFrame will be like this:

The prediction column will inform which group that entry belong to. Pay attention that the cluster numbers will allways starts with zero.

Writing the segmented image.

Spark does not provide a way to write images, so we need to implement one. Also, it is important to decide how to represent the clusters colors on the new rendered image. On this case, i decided to use the first color found from each cluster, to represent itself.

The ideia is pretty simple. For each cluster, we filter the dataset using the cluster number and get the first entry found, we add the cluster number as the key in a map and the color channels as the value in a Scala Tuple (https://docs.scala-lang.org/tour/tuples.html).

Now it is possible to finally write the final image. For that, we gonna use the class javax.imageio.ImageIO. First, we select the w, h and prediction columns and transform the DataFrame into an array, by using the collect() function. Then, we create a Buffered Image, passing the width, height and the image type. Finally, just iterate over the array, writing into the buffered image the pixel position and the respective colored image. Notice that the BufferedImage class only has the setRGB method, so we must change the order of the colors when using the class java.awt.Color.

Testing the implementation

Ok, now we can test the solution with the images presented before. But first, we must decide how many clusters we want to have for the segmentation? Looking on the image below, i will use K =5, based on the number of distinct parts that i can observe on naked eye (the red dress, the skin, the ocean, the sky and the ground).

This is the result for K = 5:

We can see that the dress and skin were segmented almost perfectly. Some parts were put into the same color of the rocks, because their spectrum are closer. The sky also is almost an unique segment. Following the next image, i will consider K = 5 also.

This is the result:

The three persons on the beach were classified in the same cluster of the montain and part of the beach, because due to the shadows, they color spectrum are similar.

Here are other runs with different K’s. Pay attention that by increasing K, we get closer to the original image. This can be useful to compress images. Try to run with K = 128 and K = 256.

K=3,K=4 and K=6

K=3,K=4 and K=6

Finding the best value of K

How can we determine the best value of K for a certain image? Previously, i said that i used my own perception to use a number, but we cannot rely on human perception, this is Data Science!

Here it is how we will proceed: we will define a range for K, like from 2 to 20. Then, for each K value, we will compute the cost of the K-Means function, and plot the cost on a graph. The value which is minimized with little

return gained thereafter, will be the best choice. For plotting, i will use the vegas-viz plot lib for scala (https://github.com/vegas-viz/Vegas). Lets include the following dependencies in the pom.xml:

<dependency> <groupId>org.vegas-viz</groupId> <artifactId>vegas_2.11</artifactId> <version>0.3.11</version> </dependency> <dependency> <groupId>org.vegas-viz</groupId> <artifactId>vegas-spark_2.11</artifactId> <version>0.3.11</version> </dependency>

We will us eMLlib ‘s KMeans() estimator in the Spark DataFrame containing our feature vectors, iterating over values of k in the range(2, 20). Then we plot the results on a html page (you can plot on a popup window, check vegas-viz documentation, but you will need the javafx package, which is not included on openjdk). Observe that we do not need to run the estimator in the entire DataFrame, taking a sample of 10% only, to save runtime execution.

So, as we can see on the results, after K=14 for the beach image does not change the cost too much and K=12 for the woman image. The new rendered images are below

K = 12

K = 14

We can see that the woman skin and dress are still in one cluster for each, which can be very usefull for detecting the person on the picture. However, the sky is not just an unique cluster. In the beach image, we can see that the three persons can be identified distinctly from the reflection on the sand.

Conclusion

I hope that i contributed to your learning on K-Means algorithm, Apache Spark and ML. Spark is a framework designed for batch processing. It extends the MapReduce model to efficiently use it for more types of computations which includes interactive queries and stream processing.

Although Spark has many ML algorithms implementations ready to use, it does not have nothing related to Computer Vision, except for the ImageSchema, which made possible load images into DataFrames. While Python developers can find a wide range of libraries and documentations about different approaches of different problems and Mix frameworks for it, i sense that the Scala/Java community is needy about more content regarding data science. This was my motivation for this development.

You can check the full code here: https://github.com/gsjunior86/SIS. If you have any questions, doubts about it, please write me.