{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Text and Image Joined representations using COLA\n\nI trained a custom model to produce image and text embeddings using COLA: https://arxiv.org/abs/2010.10915\n\nThose embeddings can be used as inputs to the recommendation system.\n\nThis is the inference script where we probe the embeddings to better understand what they contain.\n\nThe first part is a visualization of the embeddings using TSNE and the second part is using those embeddings to find the most similar images to a query image.\n\n\nThis notebook was made by @CVxTz, if you like it, you might also like some similar projects here: https://medium.com/@CVxTz","metadata":{}},{"cell_type":"code","source":"\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nfrom pathlib import Path\n\nimport os\nfor dirname, _, filenames in os.walk('../input/image-text-embeddings'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n        \nBASE_PATH = Path(\"/kaggle/input/h-and-m-personalized-fashion-recommendations/\")\nMODEL_PATH = Path(\"/kaggle/input/image-text-embeddings/ssl_resnet18_1337.ckpt\")\nTOKENIZER_PATH = Path(\"/kaggle/input/image-text-embeddings/tokenizer.json\")\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Text Utils","metadata":{}},{"cell_type":"code","source":"from tokenizers import Tokenizer\n\n\nTOKENIZER = Tokenizer.from_file(str(TOKENIZER_PATH))\n\n\nCLS_IDX = TOKENIZER.token_to_id(\"[CLS]\")\nPAD_IDX = TOKENIZER.token_to_id(\"[PAD]\")\nSEP_IDX = TOKENIZER.token_to_id(\"[SEP]\")\n\n\ndef tokenize(text: str):\n    raw_tokens = TOKENIZER.encode(text)\n\n    return raw_tokens.ids\n\n\ndef pad_list(\n    list_integers, context_size: int = 90, pad_val: int = PAD_IDX, mode=\"right\"\n):\n    \"\"\"\n\n    :param list_integers:\n    :param context_size:\n    :param pad_val:\n    :param mode:\n    :return:\n    \"\"\"\n\n    list_integers = list_integers[:context_size]\n\n    if len(list_integers) < context_size:\n        if mode == \"left\":\n            list_integers = [pad_val] * (\n                context_size - len(list_integers)\n            ) + list_integers\n        else:\n            list_integers = list_integers + [pad_val] * (\n                context_size - len(list_integers)\n            )\n\n    return list_integers\n","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Image utils","metadata":{}},{"cell_type":"code","source":"import random\nfrom pathlib import Path\n\nimport cv2\nimport numpy as np\n\ncv2.setNumThreads(0)\ncv2.ocl.setUseOpenCL(False)\n\nimport albumentations as A\n\nSIZE = 128\nSCALE = 255.0\n\n\nRESIZE = A.Compose(\n    [\n        A.LongestMaxSize(max_size=SIZE, p=1.0),\n        A.PadIfNeeded(min_height=SIZE, min_width=SIZE, p=1.0),\n    ]\n)\nNORMALIZE = A.Normalize(\n    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=SCALE\n)\n\ndef read_image(image_path: Path) -> np.ndarray:\n\n    bgr_image = cv2.imread(str(image_path))\n\n    rgb_image = bgr_image[:, :, ::-1]\n\n    return rgb_image\n\n\ndef resize(image: np.ndarray) -> np.ndarray:\n    reshaped = RESIZE(image=image)[\"image\"]\n\n    return reshaped\n\n\ndef normalize(image: np.ndarray) -> np.ndarray:\n    normalized = NORMALIZE(image=image)[\"image\"]\n\n    return normalized\n\n\ndef preprocess(image: np.ndarray) -> np.ndarray:\n    return normalize(resize(image))\n","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Model","metadata":{}},{"cell_type":"code","source":"import math\n\nimport torch\nimport torch.nn.functional as F\nfrom pytorch_lightning import LightningModule\nfrom torchvision import models\nfrom transformers import get_cosine_schedule_with_warmup\n\n\nclass PositionalEncoding(torch.nn.Module):\n    #  https://pytorch.org/tutorials/beginner/transformer_tutorial.html\n\n    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):\n        super().__init__()\n        self.dropout = torch.nn.Dropout(p=dropout)\n\n        position = torch.arange(max_len).unsqueeze(1)\n        div_term = torch.exp(\n            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)\n        )\n        pe = torch.zeros(1, max_len, d_model)\n        pe[0:, :, 0::2] = torch.sin(position * div_term)\n        pe[0:, :, 1::2] = torch.cos(position * div_term)\n        self.register_buffer(\"pe\", pe)\n\n        self.d_model = d_model\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        Args:\n            x: Tensor, shape [seq_len, batch_size, embedding_dim]\n        \"\"\"\n\n        x = x + self.pe[:, : x.size(1)] / math.sqrt(self.d_model)\n\n        return self.dropout(x)\n\n\nclass Cola(LightningModule):\n    def __init__(\n        self,\n        lr=0.001,\n        use_pretrained=False,\n        dropout=0.2,\n        d_model=128,\n        n_vocab=30_000,\n        smoothing=0.1,\n    ):\n        super().__init__()\n        self.dropout = dropout\n\n        self.lr = lr\n        self.d_model = d_model\n        self.n_vocab = n_vocab\n        self.smoothing = smoothing\n\n        # Vision\n        self.model = models.resnet18(pretrained=use_pretrained)\n        self.model.fc = torch.nn.Linear(self.model.fc.in_features, self.d_model)\n\n        # Text\n        self.item_embeddings = torch.nn.Embedding(self.n_vocab, self.d_model)\n        self.pos_encoder = PositionalEncoding(\n            d_model=self.d_model, dropout=self.dropout\n        )\n        encoder_layer = torch.nn.TransformerEncoderLayer(\n            d_model=self.d_model, nhead=4, dropout=self.dropout, batch_first=True\n        )\n        self.encoder = torch.nn.TransformerEncoder(encoder_layer, num_layers=4)\n\n        # Normalizations\n        self.layer_norm = torch.nn.LayerNorm(normalized_shape=self.d_model)\n        self.linear = torch.nn.Linear(self.d_model, self.d_model, bias=False)\n        self.do = torch.nn.Dropout(p=self.dropout)\n\n        self.save_hyperparameters()\n\n    def encode_image(self, x):\n        x = x.permute(0, 3, 1, 2)\n        x = self.do(self.model(x))\n        x = torch.tanh(self.layer_norm(x))\n\n        return x\n\n    def encode_text(self, x):\n        x = self.item_embeddings(x)\n        x = self.pos_encoder(x)\n        x = self.encoder(x)\n\n        return x[:, 0, :]\n\n    def forward(self, x):\n        image, text = x\n\n        encoded_image = self.encode_image(image)\n\n        encoded_image_w = self.linear(encoded_image)\n\n        encoded_text = self.encode_text(text)\n\n        return encoded_image_w, encoded_text\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Predict embeddings","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nimport torch\nfrom tqdm import tqdm\n\nfrom sklearn.manifold import TSNE","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_csv(\n    BASE_PATH / \"articles.csv\",\n    nrows=None,\n    dtype={\n        \"article_id\": str,\n    },\n)\n\ndf[\"text\"] = df.apply(\n    lambda x: \" \".join(\n        [\n            str(x[\"prod_name\"]),\n            str(x[\"product_type_name\"]),\n            str(x[\"product_group_name\"]),\n            str(x[\"graphical_appearance_name\"]),\n            str(x[\"colour_group_name\"]),\n            str(x[\"perceived_colour_value_name\"]),\n            str(x[\"index_name\"]),\n            str(x[\"section_name\"]),\n            str(x[\"detail_desc\"])\n        ]\n    ),\n    axis=1,\n)\n\ndf[\"image_path\"] = df.article_id.apply(lambda x: BASE_PATH / \"images\" / x[:3] / f\"{x}.jpg\")\n\ndf = df.sample(n=5000)\n\nmodel = Cola(lr=1e-4)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nmodel.load_state_dict(torch.load(MODEL_PATH, map_location=device)[\"state_dict\"])\n\nmodel.to(device)\n\nmodel.eval()\n\ntext_embeddings = []\nimage_embeddings = []\n\nfor image_path, text in tqdm(\n    zip(df.image_path.values, df.text.values),\n    total=len(df),\n):\n\n    if image_path.is_file():\n        image = read_image(image_path)\n    else:\n        image = np.zeros((128, 128, 3))\n\n    image = preprocess(image)\n\n    image_t = torch.from_numpy(image).unsqueeze(0)\n    image_t = image_t.to(device)\n\n    text_t = tokenize(text)\n    text_t = torch.tensor(pad_list(text_t), dtype=torch.long, device=device).unsqueeze(0)\n\n    with torch.no_grad():\n        text_embed = model.encode_text(text_t)\n        image_embed = model.encode_image(image_t)\n\n        text_embed = text_embed.squeeze().cpu().tolist()\n        image_embed = image_embed.squeeze().cpu().tolist()\n\n    text_embeddings.append(text_embed)\n    image_embeddings.append(image_embed)\n\ntext_embeddings = np.array(text_embeddings)\nimage_embeddings = np.array(image_embeddings)\n\n","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### TSNE Representations","metadata":{}},{"cell_type":"code","source":"\ntsne = TSNE(\n    n_components=2,\n    init=\"random\",\n    random_state=0,\n    learning_rate=\"auto\",\n    n_iter=300,\n)\n\nY = tsne.fit_transform(image_embeddings)\n\nfig = plt.figure(figsize=(12, 12))\n\nfor index_name in df.index_name.unique():\n    plt.scatter(Y[df.index_name == index_name, 0], Y[df.index_name == index_name, 1], label=index_name, s=3)\n\nplt.title(\"Cola Image embeddings by index_name\")\nplt.legend()\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\ntsne = TSNE(\n    n_components=2,\n    init=\"random\",\n    random_state=0,\n    learning_rate=\"auto\",\n    n_iter=300,\n)\n\nY = tsne.fit_transform(text_embeddings)\n\nfig = plt.figure(figsize=(12, 12))\n\nfor index_name in df.index_name.unique():\n    plt.scatter(Y[df.index_name == index_name, 0], Y[df.index_name == index_name, 1], label=index_name, s=3)\n\nplt.title(\"Cola Text embeddings by index_name\")\nplt.legend()\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Search","metadata":{}},{"cell_type":"code","source":"index = 10\n\nmost_similar = np.argsort(-image_embeddings @ image_embeddings[index, :].T)[:9].tolist()\n\n_, axs = plt.subplots(3, 3, figsize=(12, 12))\naxs = axs.flatten()\nfor i, ax in zip(most_similar, axs):\n    ax.imshow(read_image(df.image_path.values[i]))\n    ax.axis('off')\n    if i == index:\n        ax.title.set_text(\"Query image\")\n    else:\n        ax.title.set_text(\"Result Image\")\nplt.axis('off')\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"index = 100\n\nmost_similar = np.argsort(-image_embeddings @ image_embeddings[index, :].T)[:9].tolist()\n\n_, axs = plt.subplots(3, 3, figsize=(12, 12))\naxs = axs.flatten()\nfor i, ax in zip(most_similar, axs):\n    ax.imshow(read_image(df.image_path.values[i]))\n    ax.axis('off')\n    if i == index:\n        ax.title.set_text(\"Query image\")\n    else:\n        ax.title.set_text(\"Result Image\")\nplt.axis('off')\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"That's it folks! We can see that those custom embeddings are well suited to the domain of the images and texts (Fashion). I'll try them on a recommender system and share the results with you soon.","metadata":{}}]}