Image Recognition in Snowflake using Snowpark Python, PyTorch, Streamlit and OpenAI
A technical guide for building image recognition applications in Snowflake.
In this article, we will review how to build image recognition applications in Snowflake using Snowpark for Python, PyTorch, Streamlit and OpenAI’s DALL-E 2 — “a new AI system that can create realistic images and art from a description in natural language”.
One of the main reasons why running these types of applications in Snowflake makes sense is because the [Python]code executes right where the [image]data lives in your Snowflake account. The other advantage is that you don’t have to manually manage and install libraries like PyTorch and many more along with their dependencies.
So before we look at the code snippets, first things first though for those that are new to some of these technologies.
What is Snowpark?
It allows developers to query data and write data applications in languages other than SQL using a set of APIs and DataFrame-style programming constructs in languages like Python, Java, and Scala. These applications run on and take advantage of the same distributed computation on Snowflake’s elastic engine as your SQL workloads. Learn more about Snowpark.
What is PyTorch?
It is one of the most popular open source machine learning frameworks that also happens to be pre-installed and available for developers to use in Snowpark via Snowflake Anaconda channel. This means that you (the developers) can load pre-trained PyTorch models in Snowpark for Python without having to install the libraries manually.
What is Streamlit?
Streamlit is a pure-Python open source application framework that enables developers to quickly and easily write, share, and deploy data applications. Learn more about Streamlit.
OpenAI
Learn more about OpenAI and DALL-E 2. (It’s all everyone is talking about these days!)
Image Recognition in Snowflake
Now let’s take a look at the key differences between the two applications reviewed in this article and the ones you can build by following the step-by-step guide linked below.
- The first app let’s the user upload an image
- The second app uses OpenAI’s DALL-E 2 to generate an image based on user input in text/natural language format
Note: In both applications the Snowpark Python UDF code that uses PyTorch for image recognition running in Snowflake is exactly the same. Which is awesome!
Application 1 — Upload an image
This app uses Streamlit’s st.file_uploader() to allow the user to upload an image file. Once the file is uploaded successfully, the following code snippet converts image data from base64 to hex and stores it in a Snowflake table using a very handy Snowpark API session.write_pandas().
uploaded_file = st.file_uploader("Choose an image file", accept_multiple_files=False, label_visibility='hidden')
if uploaded_file is not None:
# Convert image base64 string into hex
bytes_data_in_hex = uploaded_file.getvalue().hex()
# Generate new image file name
file_name = 'img_' + str(uuid.uuid4())
# Write image data in Snowflake table
df = pd.DataFrame({"FILE_NAME": [file_name], "IMAGE_BYTES": [bytes_data_in_hex]})
session.write_pandas(df, "IMAGES")Application 2 — OpenAI generated image
This app uses OpenAI’s API openai.Image.create() to generate images based on the description provided by the user in the form of text /natural language — in real-time! Then, similar to the first app, the generated image data is converted from base64 into hex and that image data is stored in a Snowflake table using a very handy Snowpark API session.write_pandas().
# Retrieve OpenAI key from environment variable
openai.api_key = os.getenv("OPENAI_API_KEY")
# Add text box for entering text
text_input = st.text_input("Enter description of your favorite animal 👇")
if text_input:
response = openai.Image.create(
prompt=text_input,
n=1,
size="512x512",
response_format="b64_json"
)
# Convert image base64 string into hex
image_bytes = response['data'][0]['b64_json']
bytes_data_in_hex = base64.b64decode(image_bytes).hex()
# Generate new image file name
file_name = 'img_' + str(uuid.uuid4())
# Decode base64 image data and generate image file that can be used to display on screen
decoded_data = base64.b64decode((image_bytes))
with open(file_name, 'wb') as f:
f.write(decoded_data)
# Write image data in Snowflake table
df = pd.DataFrame({"FILE_NAME": [file_name], "IMAGE_BYTES": [bytes_data_in_hex]})
session.write_pandas(df, "IMAGES")Snowpark for Python and PyTorch
For this particular application, I am using PyTorch implementation of MobileNet V3. (Note: All necessary model files are included in my GitHub repo linked below and I’d also like to take this moment and extend a huge thank you to the authors for the research and making the pre-trained models available under MIT License.)
Ok, so once we have access to/downloaded the pre-trained model files, we need to upload them onto Snowflake (internal) stage using Snowpark API session.file.put() so that they can be added as dependencies on the Python UDF for inference. Here’s how…
session.file.put('imagenet1000_clsidx_to_labels.txt','@dash_files',overwrite=True,auto_compress=False)
session.file.put('mobilenetv3.py','@dash_files',overwrite=True,auto_compress=False)
session.file.put('mobilenetv3-large-1cd25616.pth','@dash_files',overwrite=True,auto_compress=False)And here’s the Snowpark for Python UDF code snippet that executes on the server in your Snowflake account. This is where the pre-trained model is used for image recognition in both applications.
@cachetools.cached(cache={})
def load_class_mapping(filename):
with open(filename, "r") as f:
return f.read()
@cachetools.cached(cache={})
def load_model():
import sys
import torch
from torchvision import models, transforms
import ast
from mobilenetv3 import mobilenetv3_large
IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]
model_file = import_dir + 'mobilenetv3-large-1cd25616.pth'
imgnet_class_mapping_file = import_dir + 'imagenet1000_clsidx_to_labels.txt'
IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD = ((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
transform = transforms.Compose([
transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
])
# Load the Imagenet {class: label} mapping
cls_idx = load_class_mapping(imgnet_class_mapping_file)
cls_idx = ast.literal_eval(cls_idx)
# Load pretrained image recognition model
model = mobilenetv3_large()
model.load_state_dict(torch.load(model_file))
# Configure pretrained model for inference
model.eval().requires_grad_(False)
return model, transform, cls_idx
def load_image(image_bytes_in_str):
import os
image_file = '/tmp/' + str(os.getpid())
image_bytes_in_hex = bytes.fromhex(image_bytes_in_str)
with open(image_file, 'wb') as f:
f.write(image_bytes_in_hex)
return open(image_file, 'rb')
@udf(name='image_recognition_using_bytes',session=session,replace=True,is_permanent=True,stage_location='@dash_udfs')
def image_recognition_using_bytes(image_bytes_in_str: str) -> str:
import sys
import torch
from PIL import Image
import os
model, transform, cls_idx = load_model()
img = Image.open(load_image(image_bytes_in_str))
img = transform(img).unsqueeze(0)
# Get model output and human text prediction
logits = model(img)
outp = torch.nn.functional.softmax(logits, dim=1)
_, idx = torch.topk(outp, 1)
idx.squeeze_()
predicted_label = cls_idx[idx.item()]
return f"{predicted_label}"Build And Run Applications
For more details, and to build and run these applications in your environment, please watch this video and follow this step-by-step QuickStart Guide.
That’s it for now…
Thanks for your time and I’d love to see what you build using these technologies.
Connect with me on Twitter and LinkedIn where I share demo videos, code snippets, and other interesting technical artifacts specifically around Snowpark for Python and Streamlit.
