notebooks/resnet_features.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6c245fde",
   "metadata": {},
   "source": [
    "# Building an Image Feature Database in `torchvision`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc8ce76a-398f-4064-95d8-c7f0a399d04f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install superduperdb==0.0.12\n",
    "!pip install torchvision"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df884338",
   "metadata": {},
   "source": [
    "In this example, we show how to utilize a pre-trained network from `torchvision` to produce image features. The images are automatically fetched and stored in MongoDB. We use a subset of the CoCo dataset (https://cocodataset.org/#home) to illustrate the process.\n",
    "\n",
    "Real-life use cases for creating a database of image features using a pre-trained network in `torchvision`:\n",
    "\n",
    "1. **Image Search and Retrieval:**\n",
    "   \n",
    "   - **Use Case:** Enhance image search capabilities in e-commerce platforms.\n",
    "   - **How:** Generate image features for products using a pre-trained network. Store these features in a database for efficient image retrieval, making it easier for users to find similar products.\n",
    "\n",
    "2. **Content-Based Recommendation Systems:**\n",
    "   \n",
    "   - **Use Case:** Improve content recommendations in media streaming services.\n",
    "   - **How:** Extract image features from movie or show frames. Store these features in a database to recommend content with similar visual characteristics to users based on their preferences.\n",
    "\n",
    "3. **Facial Recognition in Security Systems:**\n",
    "   \n",
    "   - **Use Case:** Strengthen facial recognition systems in security applications.\n",
    "   - **How:** Utilize a pre-trained neural network to extract facial features from images. Store these features in a database for quick and accurate identification in security and surveillance scenarios.\n",
    "\n",
    "4. **Medical Image Analysis:**\n",
    "   \n",
    "   - **Use Case:** Assist in medical diagnostics through image analysis.\n",
    "   - **How:** Extract features from medical images (X-rays, MRIs, etc.) using a pre-trained network. Store these features to aid in the development of diagnostic tools or systems for healthcare professionals.\n",
    "\n",
    "5. **Automated Image Tagging:**\n",
    "   \n",
    "   - **Use Case:** Streamline image organization in photo libraries or social media platforms.\n",
    "   - **How:** Extract features from uploaded images using a pre-trained model. Use these features to automatically generate relevant tags, making it easier for users to search and categorize their photos.\n",
    "\n",
    "These use cases demonstrate how creating a database of image features using `torchvision` can be applied across various domains to enhance functionality and improve user experiences. Guess what, all can be done with `superduperdb` like this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a329c235",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the zip file\n",
    "!curl -O https://superduperdb-public.s3.eu-west-1.amazonaws.com/valsmall2014.zip\n",
    "\n",
    "# Unzip the contents of the zip file (assuming the file is already downloaded)\n",
    "!unzip -qq valsmall2014.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c0351da",
   "metadata": {},
   "source": [
    "## Connect to datastore \n",
    "\n",
    "First, we need to establish a connection to a MongoDB datastore via SuperDuperDB. You can configure the `MongoDB_URI` based on your specific setup. \n",
    "\n",
    "Here are some examples of MongoDB URIs:\n",
    "\n",
    "* For testing (default connection): `mongomock://test`\n",
    "* Local MongoDB instance: `mongodb://localhost:27017`\n",
    "* MongoDB with authentication: `mongodb://superduper:superduper@mongodb:27017/documents`\n",
    "* MongoDB Atlas: `mongodb+srv://<username>:<password>@<atlas_cluster>/<database>`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8985ff38",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from superduperdb import superduper\n",
    "from superduperdb.backends.mongodb import Collection\n",
    "\n",
    "# Get the MongoDB URI from the environment variable or use a default value\n",
    "mongodb_uri = os.getenv(\"MONGODB_URI\",\"mongomock://test\")\n",
    "\n",
    "# SuperDuperDB, now handles your MongoDB database\n",
    "# It just super dupers your database \n",
    "db = superduper(mongodb_uri)\n",
    "\n",
    "# Specify a collection named 'coco'\n",
    "collection = Collection('coco')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7c903f2",
   "metadata": {},
   "source": [
    "Next, we include all image URIs in MongoDB. These URIs may include a mix of local file paths (`file://...`), web URLs (`http...`), and S3 URIs (`s3://...`). Once the URIs are added, SuperDuperDB automatically loads their content into MongoDB without the need for extra overhead or job definitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01d3af2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import random\n",
    "\n",
    "from superduperdb import Document as D\n",
    "from superduperdb.ext.pillow import pil_image as i\n",
    "\n",
    "# Get a list of file URIs for all JPEG images in the 'valsmall2014' directory\n",
    "uris = [f'file://{x}' for x in glob.glob('valsmall2014/*.jpg')]\n",
    "\n",
    "# Insert documents into the 'coco' collection in the MongoDB database\n",
    "db.execute(collection.insert_many([D({'img': i(uri=uri)}) for uri in uris], encoders=(i,))) # Here the image is encoded with pillow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92f3a73e",
   "metadata": {},
   "source": [
    "To confirm the correct storage of images in the `Datalayer`, we can perform a verification check."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e43e6243",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the display function from the IPython.display module\n",
    "from IPython.display import display\n",
    "\n",
    "# Define a lambda function for displaying images with resizing to avoid potential Jupyter crashes\n",
    "display_image = lambda x: display(x.resize((round(x.size[0] * 0.5), round(x.size[1] * 0.5))))\n",
    "\n",
    "# Retrieve the 'img' attribute from the result of collection.find_one() using db.execute()\n",
    "# Note: This assumes that db is an instance of a database connection wrapped with superduperdb\n",
    "x = db.execute(collection.find_one())['img'].x\n",
    "\n",
    "# Display the image using the previously defined lambda function\n",
    "display_image(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53d696ca",
   "metadata": {},
   "source": [
    "Let's build the `torch` + `torchvision` model using the `TorchModel` wrapper from SuperDuperDB. This allows for the incorporation of custom pre- and post-processing steps along with the model's forward pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43639f58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries and modules from torchvision and torch\n",
    "from torchvision import transforms\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision.models as models\n",
    "\n",
    "import warnings\n",
    "\n",
    "# Import custom modules\n",
    "from superduperdb.ext.torch import TorchModel, tensor\n",
    "\n",
    "# Define a series of image transformations using torchvision.transforms.Compose\n",
    "t = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),   #Resize the input image to 224x224 pixels (must same as here)\n",
    "    transforms.CenterCrop((224, 224)), # Perform a center crop on the resized image\n",
    "    transforms.ToTensor(), # Convert the image to a PyTorch tensor\n",
    "    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # Normalize the tensor with specified mean and standard deviation\n",
    "])\n",
    "\n",
    "# Define a preprocess function that applies the defined transformations to an input image\n",
    "def preprocess(x):\n",
    "    try:\n",
    "        return t(x)\n",
    "    except Exception as e:\n",
    "        # If an exception occurs during preprocessing, issue a warning and return a tensor of zeros\n",
    "        warnings.warn(str(e))\n",
    "        return torch.zeros(3, 224, 224)\n",
    "\n",
    "# Load the pre-trained ResNet-50 model from torchvision\n",
    "resnet50 = models.resnet50(pretrained=True)\n",
    "\n",
    "# Extract all layers of the ResNet-50 model except the last one\n",
    "modules = list(resnet50.children())[:-1]\n",
    "resnet50 = nn.Sequential(*modules)\n",
    "\n",
    "# Create a TorchModel instance with the ResNet-50 model, preprocessing function, and postprocessing lambda\n",
    "model = TorchModel(\n",
    "    identifier='resnet50',\n",
    "    preprocess=preprocess,\n",
    "    object=resnet50,\n",
    "    postprocess=lambda x: x[:, 0, 0],  # Postprocess by extracting the top-left element of the output tensor\n",
    "    encoder=tensor(torch.float, shape=(2048,)) # Specify the encoder configuration\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e34454fa",
   "metadata": {},
   "source": [
    "To ensure the correctness of the `model`, let's test it on a single data point by setting `one=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daa87aa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assuming x is an input tensor, you're making a prediction using the configured model\n",
    "# with the one=True parameter specifying that you expect a single prediction result.\n",
    "model.predict(x, one=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f48cb81a",
   "metadata": {},
   "source": [
    "Now that the model is prepared, we can apply it to the images stored in the `Datalayer`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a899c1cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assuming X is the input data, in this case, images ('img')\n",
    "prediction_results = model.predict(\n",
    "    X='img',                # Specify the input data (images)\n",
    "    db=db,                  # Provide the database connection or object\n",
    "    select=collection.find(),  # Specify the data to be used for prediction (fetch all data from the collection)\n",
    "    batch_size=10,          # Set the batch size for making predictions\n",
    "    max_chunk_size=3000,    # Set the maximum size of data chunks processed at once\n",
    "    in_memory=False,        # Indicate that the data is not loaded entirely into memory, processed in chunks\n",
    "    listen=True             # Enable listening mode, suggesting real-time or asynchronous prediction\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d67880a",
   "metadata": {},
   "source": [
    "To confirm that the features were stored in the `Datalayer`, you can examine them in the `_outputs.img.resnet50` field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21a1b617",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute find_one() to retrieve a single document from the collection. \n",
    "result = db.execute(collection.find_one())\n",
    "\n",
    "# The purpose of unpack() is to extract or process the data\n",
    "result.unpack()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}