Compression, search, interpolation, and clustering of images using machine learning

How to use image embeddings for compression, search, interpolation and clustering

Embeddings in machine learning provide a way to create a concise, lower-dimensional representation of complex, unstructured data. Embeddings are commonly employed in natural language processing to represent words or sentences as numbers.

In an earlier article, I showed how to create a concise representation (50 numbers) of 1059x1799 HRRR images. In this article, I will show you that the embedding has some nice properties, and you can take advantage of these properties to implement use cases like compression, image search, interpolation, and clustering of large image datasets.

Compression

First of all, does the embedding capture the important information in the image? Can we take an embedding and decode it back into the original image?

Well, we won’t be able to get back the original image, since we took 2 million pixels’ values and shoved them into a vector of length=50. Still, does the embedding capture the important information in the weather forecast image?

Here’s the original HRRR forecast on Sep 20, 2019 for 05:00 UTC:

Image for post
Image for post

We can obtain the embedding for the timestamp and decode it as follows (full code is on GitHub). First, we create a decoder by loading the SavedModel, finding the embedding layer and reconstructing all the subsequent layers:

import tensorflow as tf
def create_decoder(model_dir):
model = tf.keras.models.load_model(model_dir)
decoder_input = tf.keras.Input([50], name='embed_input')
embed_seen = False
x = decoder_input
for layer in model.layers:
if embed_seen:
x = layer(x)
elif layer.name == 'refc_embedding':
embed_seen = True
decoder = tf.keras.Model(decoder_input, x, name='decoder')
print(decoder.summary())
return decoder
decoder = create_decoder('gs://ai-analytics-solutions-kfpdemo/wxsearch/trained/savedmodel')

Once we have the decoder, we can pull the embedding for the time stamp from BigQuery:

SELECT *
FROM advdata.wxembed

This gives us a table like this:

Image for post
Image for post

We can then pass the “ref” values from the table above to the decoder:

import tensorflow as tf
import numpy as np
embed = tf.reshape( tf.convert_to_tensor(df['ref'].values[0],
dtype=tf.float32), [-1, 50])
outimg = decoder.predict(embed).squeeze() * 60
plt.imshow(outimg, origin='lower');

Note that TensorFlow expects to see a batch of inputs, and since we are passing in only one, I have to reshape it to be [1, 50]. Similarly, TensorFlow returns a batch of images. I squeeze it (remove the dummy dimension) before displaying it. The result?

Image for post
Image for post

As you can see, the decoded image is a blurry version of the original HRRR. The embedding does retain key information. It functions as a compression algorithm.

Search

If the embeddings are a compressed representation, will the degree of separation in embedding space translate to the degree of separation in terms of the actual forecast images?

If this is the case, it becomes easy to search for “similar” weather situations in the past to some scenario in the present. Finding analogs on the 2-million-pixel representation can be difficult because storms could be slightly offset from each other, or somewhat vary in size.

Since we have the embeddings in BigQuery, let’s use SQL to search for images that are similar to what happened on Sep 20, 2019 at 05:00 UTC:

WITH ref1 AS (
SELECT time AS ref1_time, ref1_value, ref1_offset
FROM `ai-analytics-solutions.advdata.wxembed`,
UNNEST(ref) AS ref1_value WITH OFFSET AS ref1_offset
WHERE time = '2019-09-20 05:00:00 UTC'
)
SELECT
time,
SUM( (ref1_value - ref[OFFSET(ref1_offset)]) * (ref1_value - ref[OFFSET(ref1_offset)]) ) AS sqdist
FROM ref1, `ai-analytics-solutions.advdata.wxembed`
GROUP BY 1
ORDER By sqdist ASC
LIMIT 5

Basically, we are computing the Euclidean distance between the embedding at the specified timestamp (refl1) and every other embedding, and displaying the closest matches. The result:

Image for post
Image for post

This makes a lot of sense. The image from the previous/next hour is the most similar. Then, images from +/- 2 hours and so on.

What if we want to find the most similar image that is not within +/- 1 day? Since we have only 1 year of data, we are not going to great analogs but let’s see what we get:

WITH ref1 AS (
SELECT time AS ref1_time, ref1_value, ref1_offset
FROM `ai-analytics-solutions.advdata.wxembed`,
UNNEST(ref) AS ref1_value WITH OFFSET AS ref1_offset
WHERE time = '2019-09-20 05:00:00 UTC'
)
SELECT
time,
SUM( (ref1_value - ref[OFFSET(ref1_offset)]) * (ref1_value - ref[OFFSET(ref1_offset)]) ) AS sqdist
FROM ref1, `ai-analytics-solutions.advdata.wxembed`
WHERE time NOT BETWEEN '2019-09-19' AND '2019-09-21'
GROUP BY 1
ORDER By sqdist ASC
LIMIT 5

The result is a bit surprising: Jan. 2 and July 1 are the days with the most similar weather:

Image for post
Image for post

Well, let’s take a look at the two timestamps:

Image for post
Image for post
Image for post
Image for post

We see that the Sep 20 image does fall somewhere between these two images. There is weather in Gulf Coast and upper midwest in both images. As it is in the Sep 20 image.

We would probably get more meaningful search if we had (a) more than just one year of data (b) loaded HRRR forecast images at multiple time-steps instead of just the analysis fields, and (c) used smaller tiles so as to capture mesoscale phenomena. This is left as an exercise to interested meteorology students reading this :)

Interpolation

Recall that when we looked for the images that were most similar to the image at 05:00, we got the images at 06:00 and 04:00 and then the images at 07:00 and 03:00. The distance to the next hour was on the order of sqrt(0.5) in embedding space.

Given this behavior in the search use case, a natural question to ask is whether we can use the embeddings for interpolating between weather forecasts. Can we average the embeddings at t-1 and t+1 to get the one at t=0? What’s the error?

WITH refl1 AS (
SELECT ref1_value, idx
FROM `ai-analytics-solutions.advdata.wxembed`,
UNNEST(ref) AS ref1_value WITH OFFSET AS idx
WHERE time = '2019-09-20 05:00:00 UTC'
),
...SELECT SUM( (ref2_value - (ref1_value + ref3_value)/2) * (ref2_value - (ref1_value + ref3_value)/2) ) AS sqdist
FROM refl1
JOIN refl2 USING (idx)
JOIN refl3 USING (idx)

The result? sqrt(0.1), which is much less than sqrt(0.5). In other words, the embeddings do function as a handy interpolation algorithm.

In order to use the embeddings as a useful interpolation algorithm, though, we need to represent the images by much more than 50 pixels. The information lost can not be this high. Again, this is left as an exercise to interested meteorologists.

Clustering

Given that the embeddings seem to work really well in terms of being commutative and additive, we should expect to be able to cluster the embeddings.

Let’s use the K-Means algorithm and ask for five clusters:

CREATE OR REPLACE MODEL advdata.hrrr_clusters
OPTIONS(model_type='kmeans', num_clusters=5, KMEANS_INIT_METHOD='KMEANS++')
AS
SELECT arr_to_input(ref) AS ref
FROM `ai-analytics-solutions.advdata.wxembed`

The resulting centroids form a 50-element array:

Image for post
Image for post

and we can go ahead and plot the decoded versions of the five centroids:

for cid in range(1,6):
embed = df[df['centroid_id']==cid].sort_values(by='feature')['numerical_value'].values
embed = tf.reshape( tf.convert_to_tensor(embed, dtype=tf.float32), [-1, 50])
outimg = decoder.predict(embed).squeeze() * 60
axarr[ (cid-1)//2, (cid-1)%2].imshow(outimg, origin='lower');
12.15992

Here are the resulting centroids of the 5 clusters:

Image for post
Image for post

The first one seems to be your class midwestern storm. The second one consists of widespread weather in the Chicago-Cleveland corridor and the Southeast. The third one is a strong variant of the second. The fourth is a squall line marching across the Appalachians. The fifth is clear skies in the interior, but weather on the coasts.

In order to use the clusters as a useful forecasting aid, though, you probably will want to cluster much smaller tiles, perhaps 500km x 500km tiles, not the entire CONUS.

In all five clusters, it is raining in Seattle and sunny in California.

Image for post
Image for post

Next steps:

  1. Peruse the full code on GitHub
  2. Read the two earlier articles. One is on how to convert HRRR files into TensorFlow records, and the other is on how to use autoencoders to create embeddings of the HRRR analysis images.
  3. If you will be attending the American Meteorological Society’s annual meeting in January 2021, come attend my talk.

Written by

Data Analytics & AI @ Google Cloud

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store