Harnessing Qdrant Vector Search for Medical Image Retrieval and Visualization
This article explores the potential of Qdrant, a vector search database, for medical image retrieval and visualization tasks.
The world of medical imaging is on fire! Deep learning models, especially the super-powered transformer models, are changing how we look at X-rays, MRIs, and CT scans. These models can now classify, segment, and analyze these images, finding hidden details doctors might miss. But there’s a catch: all this amazing data needs a smarter way to be found and explored. Here’s where vector databases like Qdrant come in. Forget keyword searches — these databases use something called “image embeddings.” Think of it like a fingerprint for each image, capturing its most important features. This lets doctors search for similar images based on meaning, not just keywords. Imagine finding all the scans with similar abnormalities in seconds!
This article explores how deep learning and vector search work together. We’ll see how Qdrant helps to find and visualize medical images faster and more easily. This could lead to better diagnoses, treatment plans, and ultimately, better care for patients.
Medical Segmentation Decathlon (MSD)
Let’s first get started with understanding the dataset. I am going to use the data from the MSD database. It contains a large medical image dataset for the development and evaluation of segmentation algorithms. The following image presents an overview of the different tasks of the MSD.
Installing the dependencies and setting up the environment:
!pip install qdrant-client transformers datasets torch numpy nibabel
!python -c "import monai" || pip install -q "monai-weekly[gdown, nibabel, tqdm, ignite]"# Importing the libraries
from monai.utils import first, set_determinism
from monai.transforms import (
AsDiscrete,
AsDiscreted,
EnsureChannelFirstd,
Compose,
CropForegroundd,
LoadImage,
LoadImaged,
Orientationd,
RandCropByPosNegLabeld,
SaveImaged,
ScaleIntensityRanged,
Spacingd,
Invertd,
DivisiblePadd,
RandAffined,
RandRotated,
RandGaussianNoised
)
from monai.handlers.utils import from_engine
from monai.networks.nets import UNet
from monai.networks.layers import Norm
from monai.metrics import DiceMetric
from monai.losses import DiceLoss, DiceCELoss
from monai.inferers import sliding_window_inference
from monai.data import CacheDataset, DataLoader, Dataset, decollate_batch
from monai.config import print_config
from monai.apps import download_and_extract
import torch
from torch.utils.data import ConcatDataset
import matplotlib.pyplot as plt
import tempfile
import shutil
import os
import glob
from datetime import datetime
import nibabel as nib
import matplotlib.pyplot as plt
# Loading the data
heart_data_path = '/content/drive/MyDrive/med_img_qdrant/Task02_Heart'
hippocampus_data_path = '/content/drive/MyDrive/med_img_qdrant/Task04_Hippocampus/Task04_Hippocampus'
prostate_data_path = '/content/drive/MyDrive/med_img_qdrant/Task05_Prostate'
spleen_data_path = '/content/drive/MyDrive/med_img_qdrant/Task09_Spleen/Task09_Spleen'
train_images = sorted(glob.glob(os.path.join(spleen_data_path, "imagesTr", "*.nii.gz")))
train_labels = sorted(glob.glob(os.path.join(spleen_data_path, "labelsTr", "*.nii.gz")))
data_dicts = [{"image": image_name, "label": label_name} for image_name, label_name in zip(train_images, train_labels)]
train_files, val_files = data_dicts[:-4], data_dicts[-4:] # 16 train and 4 validation
train_filesNow there are 2D or 3D NIfTI files available for different tasks in the dataset; we can load and visualize these images using NiBabel with Matplotlib Python package.
# sample a subset of training image data
load_img = [{'image': '/content/drive/MyDrive/med_img_qdrant/Task05_Prostate/imagesTr/prostate_00.nii.gz',
'label': '/content/drive/MyDrive/med_img_qdrant/Task05_Prostate/labelsTr/prostate_00.nii.gz'},
{'image': '/content/drive/MyDrive/med_img_qdrant/Task05_Prostate/imagesTr/prostate_01.nii.gz',
'label': '/content/drive/MyDrive/med_img_qdrant/Task05_Prostate/labelsTr/prostate_01.nii.gz'},]
"""
Display 3D image slices using Matplotlib subplots.
"""
def display_slices(image_data):
num_slices = image_data.shape[2]
num_rows = num_slices // 5 + 1
fig, axes = plt.subplots(num_rows, 5, figsize=(15, num_rows * 3))
for i in range(num_slices):
ax = axes[i // 5, i % 5]
ax.imshow(image_data[:, :, i, 0], cmap='gray') # Displaying the first volume, change index as needed
ax.set_title(f"Slice {i}")
ax.axis('off')
for j in range(num_slices, num_rows * 5):
axes[j // 5, j % 5].axis('off')
plt.tight_layout()
plt.show()
image_path = '/content/drive/MyDrive/med_img_qdrant/Task05_Prostate/imagesTr/prostate_00.nii.gz'
image_data = nib.load(image_path).get_fdata()
display_slices(image_data)
"""
Display 2D images and labels using Matplotlib subplots.
"""
for file_paths in load_img:
image_path = file_paths['image']
label_path = file_paths['label']
image_data = nib.load(image_path).get_fdata()
label_data = nib.load(label_path).get_fdata()
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.imshow(image_data[:, :, image_data.shape[2] // 2], cmap='gray')
plt.title('Image')
plt.axis('off')
plt.subplot(1, 2, 2)
plt.imshow(label_data[:, :, label_data.shape[2] // 2], cmap='gray')
plt.title('Label')
plt.axis('off')
plt.show()Qdrant: The Heart of Vector Spaces
In recent times, there has been a surge in the popularity of vector databases, which store and index vector embeddings for efficient retrieval and similarity searches. Qdrant emerges as a noteworthy player in this domain, offering an API service that streamlines the search for the closest high-dimensional vectors. As an open-source alternative to Pinecone, Qdrant boasts an intuitive API and delivers nearest-neighbor search capabilities with top-notch speed. What sets Qdrant apart from competitors like Elasticsearch is its support for filtering results based on additional payload associated with vectors, thus enhancing its versatility and utility.
Now, let’s delve into the process of constructing a cloud vector database cluster on Qdrant Cloud. To kickstart this endeavor, we’ll initiate a free-tier cloud cluster tailored for our experimental purposes. It’s worth noting that configurations can be fine-tuned to suit specific use cases and requirements, ensuring optimal performance and scalability.
If you are new to the Qdrant Cloud Database, check my previous post about setting up the Qdrant Cloud cluster and monitoring the vector database using the Thunder-HTTP client. Do follow the steps for building a new cluster on the qdrant cloud !
Now we first need to extract embeddings of our images, so that we can upsert them as vectors to perform various retrieval-related tasks, in the upcoming section I will describe a brief description of the 3D U-Net Model, which we will be using to get the image embeddings.
Image embedding extraction using UNET-3D
A 3D U-Net is a convolutional neural network architecture well-suited for extracting embeddings from medical images. Embeddings are compressed representations that capture the essential features of the image, allowing for further analysis tasks like classification or retrieval. Here’s how a 3D U-Net works for this purpose:
Architecture: It follows a U-shaped structure with two main parts:
- Encoder (contracting path): Processes the image, capturing features at different resolutions. This typically involves repeated stacks of convolutional layers with activation functions (like ReLU) followed by pooling operations (like max pooling) to reduce dimensionality.
- Decoder (expanding path): Reconstructs the embedding while preserving spatial information. It upsamples the captured features and combines them with corresponding features from the encoder path (via skip connections) at each resolution level. This allows the decoder to learn more precise and localized embeddings.
Now that the UNET-3D model is vastly adopted in medical image segmentation-related tasks, I am pointing to the code to train the model and extract image embeddings here.
Now that we have our vector embeddings of images in the data, let’s create some payload to distinguish meta-data information. I am going to append “image_name” and “body_organ” as payload to the vectors.
After completing the training, and restructuring the data into a dictionary, we can load this as a Pandas data frame as it looks like this:
All right!! So now we are going to create our collection in the cluster using the Qdrant Python client as:
import torch
from qdrant_client import QdrantClient
from qdrant_client.http import models
qdrant_client = QdrantClient(
url="https://00b127bd-0edc-4e80-b72f-xxxxxxxxxxxx.us-east4-0.gcp.cloud.qdrant.io:6333",
api_key="DJaLRz64xxxxxxxxxxxxxxxxxxxxxGcSfup7BlxiL_fO-waM-rfGAoQ",
) # replace url, api-key with your own :)
new_collection = "med_img_collection"
# our embedding vector size is 64
qdrant_client.recreate_collection(
collection_name=new_collection,
vectors_config=models.VectorParams(size=64, distance=models.Distance.COSINE)
)Now the essential part: we will create a payload, which will comprise metadata information that will also be inserted into the vector database along with vector embeddings.
payload = df[['image', 'body_organ']].to_dict(orient='records')Now finally we are ready to upsert the embedding vectors and their payload to our “med_img_collection”.
Hurray!! We have successfully uploaded our medical imaging data into our online vector database/ In the following section we will be looking at tools that Qdrant provides to investigate and extract insights from the vector store.
Semantic Search over Vector Embeddings
Semantic search is an approach used in vector databases and image retrieval systems to locate information or images based on their content or meaning rather than relying solely on keywords or metadata. This method revolutionizes the search experience by allowing users to describe what they’re looking for in natural language or even upload an image, which is then processed into a numerical representation known as an embedding vector. The system then leverages advanced algorithms to analyze the content of images and identify matches that closely align with the user’s query or input image.
let’s look at a few options to describe the vector collection:
# For printing count of the vector points
qdrant_client.count(
collection_name="med_img_collection",
exact=True,
)
# For scrolling through collection points
qdrant_client.scroll(
collection_name="med_img_collection",
limit=5
)Qdrant, a powerful tool in this domain, facilitates semantic search through its client.search() method, allowing users to seamlessly navigate through vast collections of data with ease and accuracy.
In the above figure, I gave an image embedding as the query vector and got back other images that are closest to the query image. The similarity score also gives us a good indication regarding the similarity of our query image and those in our database (excluding the first one which is the image itself, of course).
Conclusion
In wrapping up, this article has delved into the exciting potential of Qdrant, a unique tool for searching medical images. As medical imaging evolves with powerful technologies like transformer models, the way we analyze X-rays, MRIs, and CT scans is rapidly changing. These models can spot details that might escape the human eye. However, finding and exploring all this data efficiently is a challenge. That’s where Qdrant steps in. Instead of relying on simple keyword searches, it uses something called “image embeddings” to capture the essence of each image, making it easier to find similar ones based on their content. Just imagine being able to find scans with similar issues in a matter of seconds! By combining deep learning with innovative search techniques, Qdrant opens up new possibilities for medical professionals, potentially leading to better diagnoses and treatments for patients.
Follow me on Twitter: @sidgraph
(Note: This blogpost is in collaboration with Superteams.ai.)
