changed to pipelines

app.py

@@ -1,102 +1,97 @@
 import gradio as gr
-import numpy as np
-from ultralytics import YOLO
-from torchvision.transforms.functional import to_tensor
-from huggingface_hub import hf_hub_download
 import torch
-import albumentations as A
-from albumentations.pytorch.transforms import ToTensorV2
-import pandas as pd
+import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
+import pandas as pd
+from PIL import Image, ImageDraw
+import matplotlib.pyplot as plt
+import matplotlib
 
-from utils import *
-from models import YOLOStamp, Encoder
-
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-
-yolov8 = YOLO(hf_hub_download('stamps-labs/yolov8-finetuned', filename='best.torchscript'), task='detect')
+from pipelines.detection.yolo_v8 import Yolov8Pipeline
+from pipelines.detection.yolo_stamp import YoloStampPipeline
+from pipelines.segmentation.deeplabv3 import DeepLabv3Pipeline
+from pipelines.feature_extraction.vae import VaePipeline
+from pipelines.feature_extraction.vits8 import Vits8Pipeline
 
-yolo_stamp = YOLOStamp()
-yolo_stamp.load_state_dict(torch.load(hf_hub_download('stamps-labs/yolo-stamp', filename='state_dict.pth'), map_location='cpu'))
-yolo_stamp = yolo_stamp.to(device)
-yolo_stamp.eval()
-transform = A.Compose([
-    A.Normalize(),
-    ToTensorV2(p=1.0),
-])
+from utils import *
 
-vits8 = torch.jit.load(hf_hub_download('stamps-labs/vits8-stamp', filename='vits8stamp-torchscript.pth'), map_location='cpu')
-vits8 = vits8.to(device)
-vits8.eval()
 
-encoder = Encoder()
-encoder.load_state_dict(torch.load(hf_hub_download('stamps-labs/vae-encoder', filename='encoder.pth'), map_location='cpu'))
-encoder = encoder.to(device)
-encoder.eval()
+yolov8 = Yolov8Pipeline.from_pretrained(local_model_path='yolov8_old_backup.pt')
+yolo_stamp = YoloStampPipeline.from_pretrained('stamps-labs/yolo-stamp', 'weights.pt')
+vae = VaePipeline.from_pretrained('stamps-labs/vae-encoder', 'weights.pt')
+vits8 = Vits8Pipeline.from_pretrained('stamps-labs/vits8-stamp', 'weights.pt')
+dlv3 = DeepLabv3Pipeline.from_pretrained('stamps-labs/deeplabv3-finetuned', 'weights.pt')
 
 
-def predict(image, det_choice, emb_choice):
+def doc_predict(image, det_choice, seg_choice, emb_choice):
 
-    shape = torch.tensor(image.size)
     image = image.convert('RGB')
     
     if det_choice == 'yolov8':
-        coef = torch.hstack((shape, shape)) / 640
-        image = image.resize((640, 640))
-        boxes = yolov8(image)[0].boxes.xyxy.cpu()
-        image_with_boxes = visualize_bbox(image, boxes)
+        boxes = yolov8(image)
         
     elif det_choice == 'yolo-stamp':
-        coef = torch.hstack((shape, shape)) / 448
-        image = image.resize((448, 448))
-        image_tensor = transform(image=np.array(image))['image']
-        output = yolo_stamp(image_tensor.unsqueeze(0).to(device))
-
-        boxes = output_tensor_to_boxes(output[0].detach().cpu())
-        boxes = nonmax_suppression(boxes)
-        boxes = xywh2xyxy(torch.tensor(boxes)[:, :4])
-        image_with_boxes = visualize_bbox(image, boxes)
+        boxes = yolo_stamp(image)
     else:
         return
-    
+    image_with_boxes = visualize_bbox(image, boxes)
+
+    segmented_stamps = []
+    for box in boxes:
+        cropped_stamp = image.crop(box.tolist())
+        segmented_stamps.append(dlv3(cropped_stamp) if seg_choice else cropped_stamp)
+
+    widths, heights = zip(*(i.size for i in segmented_stamps))
+
+    total_width = sum(widths)
+    max_height = max(heights)
+
+    concatenated_stamps = Image.new('RGB', (total_width, max_height))
+
+    x_offset = 0
+    for im in segmented_stamps:
+        concatenated_stamps.paste(im, (x_offset,0))
+        x_offset += im.size[0]
 
     embeddings = []
     if emb_choice == 'vits8':
-        for box in boxes:
-            cropped_stamp = to_tensor(image.crop(box.tolist()))
-            embeddings.append(vits8(cropped_stamp.unsqueeze(0).to(device))[0].detach().cpu())
+        for stamp in segmented_stamps:
+            embeddings.append(vits8(stamp))
 
     elif emb_choice == 'vae-encoder':
-        for box in boxes:
-            cropped_stamp = to_tensor(image.crop(box.tolist()).resize((118, 118)))
-            embeddings.append(np.array(encoder(cropped_stamp.unsqueeze(0).to(device))[0][0].detach().cpu()))
+        for stamp in segmented_stamps:
+            embeddings.append(vae(stamp))
     
     embeddings = np.stack(embeddings)
 
     similarities = cosine_similarity(embeddings)
 
-    boxes = boxes * coef
     df_boxes = pd.DataFrame(boxes, columns=['x1', 'y1', 'x2', 'y2'])
     
     fig, ax = plt.subplots()
     im, cbar = heatmap(similarities, range(1, len(embeddings) + 1), range(1, len(embeddings) + 1), ax=ax,
                     cmap="YlGn", cbarlabel="Embeddings similarities")
     texts = annotate_heatmap(im, valfmt="{x:.3f}")
-    return image_with_boxes, df_boxes, embeddings, fig
+    return image_with_boxes, df_boxes, concatenated_stamps, embeddings, fig
 
 
-examples = [['./examples/1.jpg', 'yolov8', 'vits8'], ['./examples/2.jpg', 'yolov8', 'vae-encoder'], ['./examples/3.jpg', 'yolo-stamp', 'vits8']]
-inputs = [
-    gr.Image(type="pil"),
+doc_examples = [['examples/1.jpg', 'yolov8', True, 'vits8'], ['examples/2.jpg', 'yolo-stamp', False, 'vae-encoder'], ['examples/3.jpg', 'yolov8', True, 'vits8']]
+doc_inputs = [
+    gr.Image(label="Document image", type="pil"),
     gr.Dropdown(choices=['yolov8', 'yolo-stamp'], value='yolov8', label='Detection model'),
+    gr.Checkbox(label="Use segmentation model"),
     gr.Dropdown(choices=['vits8', 'vae-encoder'], value='vits8', label='Embedding model'),
 ]
-outputs = [
-    gr.Image(type="pil"),
+doc_outputs = [
+    gr.Image(label="Document with bounding boxes", type="pil"),
     gr.DataFrame(type='pandas', label="Bounding boxes"),
+    gr.Image(label="Segmented stamps", type="pil"),
     gr.DataFrame(type='numpy', label="Embeddings"),
     gr.Plot(label="Cosine Similarities")
 ]
-app = gr.Interface(predict, inputs, outputs, examples=examples)
-app.launch()
+
+with gr.Blocks() as demo:
+    with gr.Tab("Signle document"): 
+        gr.Interface(doc_predict, doc_inputs, doc_outputs, examples=doc_examples)
+
+demo.launch(inline=False)

constants.py → detection_models/yolo_stamp/constants.py

@@ -23,11 +23,3 @@ STD = (0.229, 0.224, 0.225)
 MEAN = (0.485, 0.456, 0.406)
 # box color to show the bounding box on image
 BOX_COLOR = (0, 0, 255)
-
-
-# dimenstion of image embedding 
-Z_DIM = 128
-# hidden dimensions for encoder model
-ENC_HIDDEN_DIM = 16
-# hidden dimensions for decoder model
-DEC_HIDDEN_DIM = 64

detection_models/yolo_stamp/data.py

@@ -0,0 +1,141 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+from sklearn.model_selection import train_test_split
+import albumentations as A
+from albumentations.pytorch.transforms import ToTensorV2
+from PIL import Image
+
+from pathlib import Path
+from random import randint
+
+from utils import *
+
+"""
+    Dataset class for storing stamps data.
+    
+    Arguments:
+    data -- list of dictionaries containing file_path (path to the image), box_nb (number of boxes on the image), and boxes of shape (4,)
+    image_folder -- path to folder containing images
+    transforms -- transforms from albumentations package
+"""
+class StampDataset(Dataset):
+    def __init__(
+            self, 
+            data=read_data(), 
+            image_folder=Path(IMAGE_FOLDER),
+            transforms=None):
+        self.data = data
+        self.image_folder = image_folder
+        self.transforms = transforms
+
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        image_fn = self.image_folder / item['file_path']
+        boxes = item['boxes']
+        box_nb = item['box_nb']
+        labels = torch.zeros((box_nb, 2), dtype=torch.int64)
+        labels[:, 0] = 1
+
+        img = np.array(Image.open(image_fn))
+
+        try:
+            if self.transforms:
+                sample = self.transforms(**{
+                    "image":img,
+                    "bboxes": boxes,
+                    "labels": labels,
+                })
+                img = sample['image']
+                boxes = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)
+        except:
+            return self.__getitem__(randint(0, len(self.data)-1))
+
+        target_tensor = boxes_to_tensor(boxes.type(torch.float32))
+        return img, target_tensor
+
+    def __len__(self):
+        return len(self.data)
+
+def collate_fn(batch):
+    return tuple(zip(*batch))
+
+
+def get_datasets(data_path=ANNOTATIONS_PATH, train_transforms=None, val_transforms=None):
+    """
+        Creates StampDataset objects.
+
+        Arguments: 
+        data_path -- string or Path, specifying path to annotations file
+        train_transforms -- transforms to be applied during training
+        val_transforms -- transforms to be applied during validation
+
+        Returns:
+        (train_dataset, val_dataset) -- tuple of StampDataset for training and validation
+    """
+    data = read_data(data_path)
+    if train_transforms is None:
+        train_transforms = A.Compose([
+            A.RandomCropNearBBox(max_part_shift=0.6, p=0.4),
+            A.Resize(height=448, width=448),
+            A.HorizontalFlip(p=0.5),
+            A.VerticalFlip(p=0.5),
+            # A.Affine(scale=(0.9, 1.1), translate_percent=(0.05, 0.1), rotate=(-45, 45), shear=(-30, 30), p=0.3),
+            # A.Blur(blur_limit=4, p=0.3),
+            A.Normalize(),
+            ToTensorV2(p=1.0),
+        ],
+        bbox_params={
+            "format":"coco",
+            'label_fields': ['labels']
+        })
+
+    if val_transforms is None:
+        val_transforms = A.Compose([
+            A.Resize(height=448, width=448),
+            A.Normalize(),
+            ToTensorV2(p=1.0),
+        ],
+        bbox_params={
+            "format":"coco",
+            'label_fields': ['labels']
+        })
+    train, test_data = train_test_split(data, test_size=0.1, shuffle=True)
+
+    train_data, val_data = train_test_split(train, test_size=0.2, shuffle=True)
+
+    train_dataset = StampDataset(train_data, transforms=train_transforms)
+    val_dataset = StampDataset(val_data, transforms=val_transforms)
+    test_dataset = StampDataset(test_data, transforms=val_transforms)
+
+    return train_dataset, val_dataset, test_dataset
+
+
+def get_loaders(batch_size=8, data_path=ANNOTATIONS_PATH, num_workers=0, train_transforms=None, val_transforms=None):
+    """
+        Creates StampDataset objects.
+
+        Arguments: 
+        batch_size -- integer specifying the number of images in the batch
+        data_path -- string or Path, specifying path to annotations file
+        train_transforms -- transforms to be applied during training
+        val_transforms -- transforms to be applied during validation
+
+        Returns:
+        (train_loader, val_loader) -- tuple of DataLoader for training and validation
+    """
+    train_dataset, val_dataset, _ = get_datasets(data_path)
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        collate_fn=collate_fn, drop_last=True)
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        collate_fn=collate_fn)
+    
+    return train_loader, val_loader

detection_models/yolo_stamp/loss.py

@@ -0,0 +1,52 @@
+import torch
+import torch.nn as nn
+from utils import *
+
+"""
+    Class for loss for training YOLO model.
+
+    Argmunets:
+    h_coord: weight for loss related to coordinates and shapes of box
+    h__noobj: weight for loss of predicting presence of box when it is absent.
+"""
+class YOLOLoss(nn.Module):
+    def __init__(self, h_coord=0.5, h_noobj=2., h_shape=2., h_obj=10.):
+        super().__init__()
+        self.h_coord = h_coord
+        self.h_noobj = h_noobj
+        self.h_shape = h_shape
+        self.h_obj = h_obj
+    
+    def square_error(self, output, target):
+        return (output - target) ** 2
+
+    def forward(self, output, target):
+        
+        pred_xy, pred_wh, pred_obj = yolo_head(output)
+        gt_xy, gt_wh, gt_obj = process_target(target)
+
+        pred_ul = pred_xy - 0.5 * pred_wh
+        pred_br = pred_xy + 0.5 * pred_wh
+        pred_area = pred_wh[..., 0] * pred_wh[..., 1]
+
+        gt_ul = gt_xy - 0.5 * gt_wh
+        gt_br = gt_xy + 0.5 * gt_wh
+        gt_area = gt_wh[..., 0] * gt_wh[..., 1]
+
+        intersect_ul = torch.max(pred_ul, gt_ul)
+        intersect_br = torch.min(pred_br, gt_br)
+        intersect_wh = intersect_br - intersect_ul
+        intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
+
+        iou = intersect_area / (pred_area + gt_area - intersect_area)
+        max_iou = torch.max(iou, dim=3, keepdim=True)[0]
+        best_box_index = torch.unsqueeze(torch.eq(iou, max_iou).float(), dim=-1)
+        gt_box_conf = best_box_index * gt_obj
+
+        xy_loss = (self.square_error(pred_xy, gt_xy) * gt_box_conf).sum()
+        wh_loss = (self.square_error(pred_wh, gt_wh) * gt_box_conf).sum()
+        obj_loss = (self.square_error(pred_obj, gt_obj) * gt_box_conf).sum()
+        noobj_loss = (self.square_error(pred_obj, gt_obj) * (1 - gt_box_conf)).sum()
+
+        total_loss = self.h_coord * xy_loss + self.h_shape * wh_loss + self.h_obj * obj_loss + self.h_noobj * noobj_loss
+        return total_loss

detection_models/yolo_stamp/model.py

@@ -0,0 +1,80 @@
+import torch
+import torch.nn as nn
+
+from .constants import *
+
+"""
+    Class for custom activation.
+"""
+class SymReLU(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super().__init__()
+        self.inplace = inplace
+
+    def forward(self, input):
+        return torch.min(torch.max(input, -torch.ones_like(input)), torch.ones_like(input))
+
+    def extra_repr(self) -> str:
+        inplace_str = 'inplace=True' if self.inplace else ''
+        return inplace_str
+
+
+"""
+    Class implementing YOLO-Stamp architecture described in https://link.springer.com/article/10.1134/S1054661822040046.
+"""
+class YOLOStamp(nn.Module):
+    def __init__(
+            self,
+            anchors=ANCHORS,
+            in_channels=3,
+    ):
+        super().__init__()
+        
+        self.register_buffer('anchors', torch.tensor(anchors))
+
+        self.act = SymReLU()
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.norm1 = nn.BatchNorm2d(num_features=8)
+        self.conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.norm2 = nn.BatchNorm2d(num_features=16)
+        self.conv3 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.norm3 = nn.BatchNorm2d(num_features=16)
+        self.conv4 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.norm4 = nn.BatchNorm2d(num_features=16)
+        self.conv5 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.norm5 = nn.BatchNorm2d(num_features=16)
+        self.conv6 = nn.Conv2d(in_channels=16, out_channels=24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.norm6 = nn.BatchNorm2d(num_features=24)
+        self.conv7 = nn.Conv2d(in_channels=24, out_channels=24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.norm7 = nn.BatchNorm2d(num_features=24)
+        self.conv8 = nn.Conv2d(in_channels=24, out_channels=48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.norm8 = nn.BatchNorm2d(num_features=48)
+        self.conv9 = nn.Conv2d(in_channels=48, out_channels=48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.norm9 = nn.BatchNorm2d(num_features=48)
+        self.conv10 = nn.Conv2d(in_channels=48, out_channels=48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.norm10 = nn.BatchNorm2d(num_features=48)
+        self.conv11 = nn.Conv2d(in_channels=48, out_channels=64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.norm11 = nn.BatchNorm2d(num_features=64)
+        self.conv12 = nn.Conv2d(in_channels=64, out_channels=256, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0))
+        self.norm12 = nn.BatchNorm2d(num_features=256)
+        self.conv13 = nn.Conv2d(in_channels=256, out_channels=len(anchors) * 5, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0))
+    
+    def forward(self, x, head=True):
+        x = x.type(self.conv1.weight.dtype)
+        x = self.act(self.pool(self.norm1(self.conv1(x))))
+        x = self.act(self.pool(self.norm2(self.conv2(x))))
+        x = self.act(self.pool(self.norm3(self.conv3(x))))
+        x = self.act(self.pool(self.norm4(self.conv4(x))))
+        x = self.act(self.pool(self.norm5(self.conv5(x))))
+        x = self.act(self.norm6(self.conv6(x)))
+        x = self.act(self.norm7(self.conv7(x)))
+        x = self.act(self.pool(self.norm8(self.conv8(x))))
+        x = self.act(self.norm9(self.conv9(x)))
+        x = self.act(self.norm10(self.conv10(x)))
+        x = self.act(self.norm11(self.conv11(x)))
+        x = self.act(self.norm12(self.conv12(x)))
+        x = self.conv13(x)
+        nb, _, nh, nw= x.shape
+        x = x.permute(0, 2, 3, 1).view(nb, nh, nw, self.anchors.shape[0], 5)
+        return x

detection_models/yolo_stamp/train.ipynb

@@ -0,0 +1,185 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from model import *\n",
+    "from loss import *\n",
+    "from data import *\n",
+    "from torch import optim\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import pytorch_lightning as pl\n",
+    "from torchmetrics.detection import MeanAveragePrecision\n",
+    "from pytorch_lightning.loggers import TensorBoardLogger"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "_, _, test_dataset = get_datasets()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class LitModel(pl.LightningModule):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.model = YOLOStamp()\n",
+    "        self.criterion = YOLOLoss()\n",
+    "        self.val_map = MeanAveragePrecision(box_format='xywh', iou_type='bbox')\n",
+    "    \n",
+    "    def forward(self, x):\n",
+    "        return self.model(x)\n",
+    "\n",
+    "    def configure_optimizers(self):\n",
+    "        optimizer = optim.AdamW(self.parameters(), lr=1e-3)\n",
+    "        # return optimizer\n",
+    "        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, 1000)\n",
+    "        return {\"optimizer\": optimizer, \"lr_scheduler\": scheduler}\n",
+    "\n",
+    "    def training_step(self, batch, batch_idx):\n",
+    "        images, targets = batch\n",
+    "        tensor_images = torch.stack(images)\n",
+    "        tensor_targets = torch.stack(targets)\n",
+    "        output = self.model(tensor_images)\n",
+    "        loss = self.criterion(output, tensor_targets)\n",
+    "        self.log(\"train_loss\", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)\n",
+    "        return loss\n",
+    "\n",
+    "    def validation_step(self, batch, batch_idx):\n",
+    "        images, targets = batch\n",
+    "        tensor_images = torch.stack(images)\n",
+    "        tensor_targets = torch.stack(targets)\n",
+    "        output = self.model(tensor_images)\n",
+    "        loss = self.criterion(output, tensor_targets)\n",
+    "        self.log(\"val_loss\", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)\n",
+    "\n",
+    "        for i in range(len(images)):\n",
+    "            boxes = output_tensor_to_boxes(output[i].detach().cpu())\n",
+    "            boxes = nonmax_suppression(boxes)\n",
+    "            target = target_tensor_to_boxes(targets[i])[::BOX]\n",
+    "            if not boxes:\n",
+    "                boxes = torch.zeros((1, 5))\n",
+    "            preds = [\n",
+    "                dict(\n",
+    "                    boxes=torch.tensor(boxes)[:, :4].clone().detach(),\n",
+    "                    scores=torch.tensor(boxes)[:, 4].clone().detach(),\n",
+    "                    labels=torch.zeros(len(boxes)),\n",
+    "                )\n",
+    "            ]\n",
+    "            target = [\n",
+    "                dict(\n",
+    "                    boxes=torch.tensor(target),\n",
+    "                    labels=torch.zeros(len(target)),\n",
+    "                )\n",
+    "            ]\n",
+    "            self.val_map.update(preds, target)\n",
+    "    \n",
+    "    def on_validation_epoch_end(self):\n",
+    "        mAPs = {\"val_\" + k: v for k, v in self.val_map.compute().items()}\n",
+    "        mAPs_per_class = mAPs.pop(\"val_map_per_class\")\n",
+    "        mARs_per_class = mAPs.pop(\"val_mar_100_per_class\")\n",
+    "        self.log_dict(mAPs)\n",
+    "        self.val_map.reset()\n",
+    "\n",
+    "        image = test_dataset[randint(0, len(test_dataset) - 1)][0].to(self.device)\n",
+    "        output = self.model(image.unsqueeze(0))\n",
+    "        boxes = output_tensor_to_boxes(output[0].detach().cpu())\n",
+    "        boxes = nonmax_suppression(boxes)\n",
+    "        img = image.permute(1, 2, 0).cpu().numpy()\n",
+    "        img = visualize_bbox(img.copy(), boxes=boxes)\n",
+    "        img = (255. * (img * np.array(STD) + np.array(MEAN))).astype(np.uint8)\n",
+    "        \n",
+    "        self.logger.experiment.add_image(\"detected boxes\", torch.tensor(img).permute(2, 0, 1), self.current_epoch)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "litmodel = LitModel()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "logger = TensorBoardLogger(\"detection_logs\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "epochs = 100"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_loader, val_loader = get_loaders(batch_size=8)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = pl.Trainer(accelerator=\"auto\", max_epochs=epochs, logger=logger)\n",
+    "trainer.fit(model=litmodel, train_dataloaders=train_loader, val_dataloaders=val_loader)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%tensorboard"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.9.0"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

detection_models/yolo_stamp/utils.py

@@ -0,0 +1,275 @@
+import torch
+import cv2
+import pandas as pd
+import numpy as np
+from pathlib import Path
+import matplotlib.pyplot as plt
+from .constants import *
+
+
+def output_tensor_to_boxes(boxes_tensor):
+    """
+        Converts the YOLO output tensor to list of boxes with probabilites.
+
+        Arguments:
+        boxes_tensor -- tensor of shape (S, S, BOX, 5)
+
+        Returns:
+        boxes -- list of shape (None, 5)
+
+        Note: "None" is here because you don't know the exact number of selected boxes, as it depends on the threshold. 
+        For example, the actual output size of scores would be (10, 5) if there are 10 boxes
+    """
+    cell_w, cell_h = W/S, H/S
+    boxes = []
+    
+    for i in range(S):
+        for j in range(S):
+            for b in range(BOX):
+                anchor_wh = torch.tensor(ANCHORS[b])
+                data = boxes_tensor[i,j,b]
+                xy = torch.sigmoid(data[:2])
+                wh = torch.exp(data[2:4])*anchor_wh
+                obj_prob = torch.sigmoid(data[4])
+                
+                if obj_prob > OUTPUT_THRESH:
+                    x_center, y_center, w, h = xy[0], xy[1], wh[0], wh[1]
+                    x, y = x_center+j-w/2, y_center+i-h/2
+                    x,y,w,h = x*cell_w, y*cell_h, w*cell_w, h*cell_h
+                    box = [x,y,w,h, obj_prob]
+                    boxes.append(box)
+    return boxes
+
+
+def plot_img(img, size=(7,7)):
+    plt.figure(figsize=size)
+    plt.imshow(img)
+    plt.show()
+
+
+def plot_normalized_img(img, std=STD, mean=MEAN, size=(7,7)):
+    mean = mean if isinstance(mean, np.ndarray) else np.array(mean)
+    std = std if isinstance(std, np.ndarray) else np.array(std)
+    plt.figure(figsize=size)
+    plt.imshow((255. * (img * std + mean)).astype(np.uint))
+    plt.show()
+    
+
+def visualize_bbox(img, boxes, thickness=2, color=BOX_COLOR, draw_center=True):
+    """
+        Draws boxes on the given image.
+
+        Arguments:
+        img -- torch.Tensor of shape (3, W, H) or numpy.ndarray of shape (W, H, 3)
+        boxes -- list of shape (None, 5)
+        thickness -- number specifying the thickness of box border
+        color -- RGB tuple of shape (3,) specifying the color of boxes
+        draw_center -- boolean specifying whether to draw center or not
+
+        Returns:
+        img_copy -- numpy.ndarray of shape(W, H, 3) containing image with bouning boxes
+    """
+    img_copy = img.cpu().permute(1,2,0).numpy() if isinstance(img, torch.Tensor) else img.copy()
+    for box in boxes:
+        x,y,w,h = int(box[0]), int(box[1]), int(box[2]), int(box[3])
+        img_copy = cv2.rectangle(
+            img_copy,
+            (x,y),(x+w, y+h),
+            color, thickness)
+        if draw_center:
+            center = (x+w//2, y+h//2)
+            img_copy = cv2.circle(img_copy, center=center, radius=3, color=(0,255,0), thickness=2)
+    return img_copy
+
+
+def read_data(annotations=Path(ANNOTATIONS_PATH)):
+    """
+        Reads annotations data from .csv file. Must contain columns: image_name, bbox_x, bbox_y, bbox_width, bbox_height.
+
+        Arguments:
+        annotations_path -- string or Path specifying path of annotations file
+
+        Returns:
+        data -- list of dictionaries containing path, number of boxes and boxes itself
+    """
+    data = []
+
+    boxes = pd.read_csv(annotations)
+    image_names = boxes['image_name'].unique()
+
+    for image_name in image_names:
+        cur_boxes = boxes[boxes['image_name'] == image_name]
+        img_data = {
+            'file_path': image_name,
+            'box_nb': len(cur_boxes),
+            'boxes': []}
+        stamp_nb = img_data['box_nb']
+        if stamp_nb <= STAMP_NB_MAX:
+            img_data['boxes'] = cur_boxes[['bbox_x', 'bbox_y','bbox_width','bbox_height']].values
+        data.append(img_data)
+    return data
+
+def xywh2xyxy(x):
+    """
+        Converts xywh format to xyxy
+
+        Arguments:
+        x -- torch.Tensor or np.array (xywh format)
+
+        Returns:
+        y -- torch.Tensor or np.array (xyxy)
+    """
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[..., 0] = x[..., 0]
+    y[..., 1] = x[..., 1]
+    y[..., 2] = x[..., 0] + x[..., 2] 
+    y[..., 3] = x[..., 1] + x[..., 3] 
+    return y
+
+def boxes_to_tensor(boxes):
+    """
+        Convert list of boxes (and labels) to tensor format
+        
+        Arguments:
+        boxes -- list of boxes
+
+        Returns:
+        boxes_tensor -- tensor of shape (S, S, BOX, 5)
+    """
+    boxes_tensor = torch.zeros((S, S, BOX, 5))
+    cell_w, cell_h = W/S, H/S
+    for i, box in enumerate(boxes):
+        x, y, w, h = box
+        # normalize xywh with cell_size
+        x, y, w, h = x / cell_w, y / cell_h, w / cell_w, h / cell_h
+        center_x, center_y = x + w / 2, y + h / 2
+        grid_x = int(np.floor(center_x))
+        grid_y = int(np.floor(center_y))
+        
+        if grid_x < S and grid_y < S:
+            boxes_tensor[grid_y, grid_x, :, 0:4] = torch.tensor(BOX * [[center_x - grid_x, center_y - grid_y, w, h]])
+            boxes_tensor[grid_y, grid_x, :, 4]  = torch.tensor(BOX * [1.])
+    return boxes_tensor
+
+
+def target_tensor_to_boxes(boxes_tensor, output_threshold=OUTPUT_THRESH):
+    """
+        Recover target tensor (tensor output of dataset) to bboxes.
+        Arguments:
+            boxes_tensor -- tensor of shape (S, S, BOX, 5)
+        Returns:
+            boxes -- list of boxes, each box is [x, y, w, h]
+    """
+    cell_w, cell_h = W/S, H/S
+    boxes = []
+    for i in range(S):
+        for j in range(S):
+            for b in range(BOX):
+                data = boxes_tensor[i,j,b]
+                x_center,y_center, w, h, obj_prob = data[0], data[1], data[2], data[3], data[4]
+                if obj_prob > output_threshold:
+                    x, y = x_center+j-w/2, y_center+i-h/2
+                    x,y,w,h = x*cell_w, y*cell_h, w*cell_w, h*cell_h
+                    box = [x,y,w,h]
+                    boxes.append(box)
+    return boxes    
+
+
+def overlap(interval_1, interval_2):
+    """
+        Calculates length of overlap between two intervals.
+
+        Arguments:
+        interval_1 -- list or tuple of shape (2,) containing endpoints of the first interval
+        interval_2 -- list or tuple of shape (2, 2) containing endpoints of the second interval
+
+        Returns:
+        overlap -- length of overlap
+    """
+    x1, x2 = interval_1
+    x3, x4 = interval_2
+    if x3 < x1:
+        if x4 < x1:
+            return 0
+        else:
+            return min(x2,x4) - x1
+    else:
+        if x2 < x3:
+            return 0
+        else:
+            return min(x2,x4) - x3
+
+
+def compute_iou(box1, box2):
+    """
+        Compute IOU between box1 and box2.
+
+        Argmunets:
+        box1 -- list of shape (5, ). Represents the first box
+        box2 -- list of shape (5, ). Represents the second box
+        Each box is [x, y, w, h, prob]
+
+        Returns:
+        iou -- intersection over union score between two boxes
+    """
+    x1,y1,w1,h1 = box1[0], box1[1], box1[2], box1[3]
+    x2,y2,w2,h2 = box2[0], box2[1], box2[2], box2[3]
+
+    area1, area2 = w1*h1, w2*h2
+    intersect_w = overlap((x1,x1+w1), (x2,x2+w2))
+    intersect_h = overlap((y1,y1+h1), (y2,y2+w2))
+    if intersect_w == w1 and intersect_h == h1 or intersect_w == w2 and intersect_h == h2:
+        return 1.
+    intersect_area = intersect_w*intersect_h
+    iou = intersect_area/(area1 + area2 - intersect_area)
+    return iou
+
+
+def nonmax_suppression(boxes, iou_thresh = IOU_THRESH):
+    """
+        Removes ovelap bboxes
+
+        Arguments:
+        boxes -- list of shape (None, 5)
+        iou_thresh -- maximal value of iou when boxes are considered different
+        Each box is [x, y, w, h, prob]
+
+        Returns:
+        boxes -- list of shape (None, 5) with removed overlapping boxes
+    """
+    boxes = sorted(boxes, key=lambda x: x[4], reverse=True)
+    for i, current_box in enumerate(boxes):
+        if current_box[4] <= 0:
+            continue
+        for j in range(i+1, len(boxes)):
+            iou = compute_iou(current_box, boxes[j])
+            if iou > iou_thresh:
+                boxes[j][4] = 0
+    boxes = [box for box in boxes if box[4] > 0]
+    return boxes
+
+
+    
+def yolo_head(yolo_output):
+    """
+        Converts a yolo output tensor to separate tensors of coordinates, shapes and probabilities.
+
+        Arguments:
+        yolo_output -- tensor of shape (batch_size, S, S, BOX, 5)
+
+        Returns:
+        xy -- tensor of shape (batch_size, S, S, BOX, 2) containing coordinates of centers of found boxes for each anchor in each grid cell
+        wh -- tensor of shape (batch_size, S, S, BOX, 2) containing width and height of found boxes for each anchor in each grid cell
+        prob -- tensor of shape (batch_size, S, S, BOX, 1) containing the probability of presence of boxes for each anchor in each grid cell
+    """
+    xy = torch.sigmoid(yolo_output[..., 0:2])
+    anchors_wh = torch.tensor(ANCHORS, device=yolo_output.device).view(1, 1, 1, len(ANCHORS), 2)
+    wh = torch.exp(yolo_output[..., 2:4]) * anchors_wh
+    prob = torch.sigmoid(yolo_output[..., 4:5])
+    return xy, wh, prob
+
+def process_target(target):
+    xy = target[..., 0:2]
+    wh = target[..., 2:4]
+    prob = target[..., 4:5]
+    return xy, wh, prob

detection_models/yolov8/train.ipynb

@@ -0,0 +1,144 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "HOME = os.getcwd()\n",
+    "print(HOME)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Pip install method (recommended)\n",
+    "\n",
+    "%pip install ultralytics==8.0.20\n",
+    "\n",
+    "from IPython import display\n",
+    "display.clear_output()\n",
+    "\n",
+    "import ultralytics\n",
+    "ultralytics.checks()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ultralytics import YOLO\n",
+    "\n",
+    "from IPython.display import display, Image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!mkdir {HOME}/datasets\n",
+    "%cd {HOME}/datasets\n",
+    "\n",
+    "%pip install roboflow --quiet\n",
+    "\n",
+    "from roboflow import Roboflow\n",
+    "rf = Roboflow(api_key=\"YOUR_API_KEY\")\n",
+    "project = rf.workspace(\"WORKSPACE\").project(\"PROJECT\")\n",
+    "dataset = project.version(1).download(\"yolov8\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%cd {HOME}\n",
+    "\n",
+    "!yolo task=detect mode=train model=yolov8s.pt data={dataset.location}/data.yaml epochs=25 imgsz=800 plots=True"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%cd {HOME}\n",
+    "Image(filename=f'{HOME}/runs/detect/train/confusion_matrix.png', width=600)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%cd {HOME}\n",
+    "Image(filename=f'{HOME}/runs/detect/train/results.png', width=600)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%cd {HOME}\n",
+    "Image(filename=f'{HOME}/runs/detect/train/val_batch0_pred.jpg', width=600)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%cd {HOME}\n",
+    "\n",
+    "!yolo task=detect mode=val model={HOME}/runs/detect/train/weights/best.pt data={dataset.location}/data.yaml"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%cd {HOME}\n",
+    "!yolo task=detect mode=predict model={HOME}/runs/detect/train/weights/best.pt conf=0.25 source={dataset.location}/test/images save=True"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import glob\n",
+    "from IPython.display import Image, display\n",
+    "\n",
+    "for image_path in glob.glob(f'{HOME}/runs/detect/predict3/*.jpg')[:3]:\n",
+    "      display(Image(filename=image_path, width=600))\n",
+    "      print(\"\\n\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

embedding_models/vae/constants.py

@@ -0,0 +1,6 @@
+# dimenstion of image embedding 
+Z_DIM = 128
+# hidden dimensions for encoder model
+ENC_HIDDEN_DIM = 16
+# hidden dimensions for decoder model
+DEC_HIDDEN_DIM = 64

embedding_models/vae/losses.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn as nn
+from torch.distributions.kl import kl_divergence
+from torch.distributions.normal import Normal
+from torch.nn.functional import relu
+
+
+
+class BatchHardTripletLoss(nn.Module):
+    def __init__(self, margin=1., squared=False, agg='sum'):
+        """
+        Initalize the loss function with a margin parameter, whether or not to consider
+        squared Euclidean distance and how to aggregate the loss in a batch
+        """
+        super().__init__()
+        self.margin = margin
+        self.squared = squared
+        self.agg = agg
+        self.eps = 1e-8
+
+    def get_pairwise_distances(self, embeddings):
+        """
+        Computing Euclidean distance for all possible pairs of embeddings.
+        """
+        ab = embeddings.mm(embeddings.t())
+        a_squared = ab.diag().unsqueeze(1)
+        b_squared = ab.diag().unsqueeze(0)
+        distances = a_squared - 2 * ab + b_squared
+        distances = relu(distances)
+
+        if not self.squared:
+            distances = torch.sqrt(distances + self.eps)
+
+        return distances
+
+    def hardest_triplet_mining(self, dist_mat, labels):
+        
+        assert len(dist_mat.size()) == 2
+        assert dist_mat.size(0) == dist_mat.size(1)
+        N = dist_mat.size(0)
+        
+        is_pos = labels.expand(N, N).eq(labels.expand(N, N).t())
+        is_neg = labels.expand(N, N).ne(labels.expand(N, N).t())
+
+        dist_ap, relative_p_inds = torch.max(
+            (dist_mat * is_pos), 1, keepdim=True)
+        
+        dist_an, relative_n_inds = torch.min(
+            (dist_mat * is_neg), 1, keepdim=True)
+
+        return dist_ap, dist_an
+
+    def forward(self, embeddings, labels):
+        
+        distances = self.get_pairwise_distances(embeddings)
+        dist_ap, dist_an = self.hardest_triplet_mining(distances, labels)
+
+        triplet_loss = relu(dist_ap - dist_an + self.margin).sum()
+        return triplet_loss
+
+
+class VAELoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.reconstruction_loss = nn.BCELoss(reduction='sum')
+    
+    def kl_divergence_loss(self, q_dist):
+        return kl_divergence(
+            q_dist, Normal(torch.zeros_like(q_dist.mean), torch.ones_like(q_dist.stddev))
+        ).sum(-1)
+    
+    
+    def forward(self, output, target, encoding):
+        loss = self.kl_divergence_loss(encoding).sum() + self.reconstruction_loss(output, target)
+        return loss
+
+

embedding_models/vae/model.py

@@ -0,0 +1,147 @@
+import torch.nn as nn
+from torch.distributions.normal import Normal
+
+from .constants import *
+
+
+class Encoder(nn.Module):
+    '''
+    Encoder Class
+    Values:
+    im_chan: the number of channels of the output image, a scalar
+    hidden_dim: the inner dimension, a scalar
+    '''
+
+    def __init__(self, im_chan=3, output_chan=Z_DIM, hidden_dim=ENC_HIDDEN_DIM):
+        super(Encoder, self).__init__()
+        self.z_dim = output_chan
+        self.disc = nn.Sequential(
+            self.make_disc_block(im_chan, hidden_dim),
+            self.make_disc_block(hidden_dim, hidden_dim * 2),
+            self.make_disc_block(hidden_dim * 2, hidden_dim * 4),
+            self.make_disc_block(hidden_dim * 4, hidden_dim * 8),
+            self.make_disc_block(hidden_dim * 8, output_chan * 2, final_layer=True),
+        )
+
+    def make_disc_block(self, input_channels, output_channels, kernel_size=4, stride=2, final_layer=False):
+        '''
+        Function to return a sequence of operations corresponding to a encoder block of the VAE, 
+        corresponding to a convolution, a batchnorm (except for in the last layer), and an activation
+        Parameters:
+        input_channels: how many channels the input feature representation has
+        output_channels: how many channels the output feature representation should have
+        kernel_size: the size of each convolutional filter, equivalent to (kernel_size, kernel_size)
+        stride: the stride of the convolution
+        final_layer: whether we're on the final layer (affects activation and batchnorm)
+        '''        
+        if not final_layer:
+            return nn.Sequential(
+                nn.Conv2d(input_channels, output_channels, kernel_size, stride),
+                nn.BatchNorm2d(output_channels),
+                nn.LeakyReLU(0.2, inplace=True),
+            )
+        else:
+            return nn.Sequential(
+                nn.Conv2d(input_channels, output_channels, kernel_size, stride),
+            )
+
+    def forward(self, image):
+        '''
+        Function for completing a forward pass of the Encoder: Given an image tensor, 
+        returns a 1-dimension tensor representing fake/real.
+        Parameters:
+        image: a flattened image tensor with dimension (im_dim)
+        '''
+        disc_pred = self.disc(image)
+        encoding = disc_pred.view(len(disc_pred), -1)
+        # The stddev output is treated as the log of the variance of the normal 
+        # distribution by convention and for numerical stability
+        return encoding[:, :self.z_dim], encoding[:, self.z_dim:].exp()
+
+
+class Decoder(nn.Module):
+    '''
+    Decoder Class
+    Values:
+    z_dim: the dimension of the noise vector, a scalar
+    im_chan: the number of channels of the output image, a scalar
+    hidden_dim: the inner dimension, a scalar
+    '''
+    
+    def __init__(self, z_dim=Z_DIM, im_chan=3, hidden_dim=DEC_HIDDEN_DIM):
+        super(Decoder, self).__init__()
+        self.z_dim = z_dim
+        self.gen = nn.Sequential(
+            self.make_gen_block(z_dim, hidden_dim * 16),
+            self.make_gen_block(hidden_dim * 16, hidden_dim * 8, kernel_size=4, stride=1),
+            self.make_gen_block(hidden_dim * 8, hidden_dim * 4),
+            self.make_gen_block(hidden_dim * 4, hidden_dim * 2, kernel_size=4),
+            self.make_gen_block(hidden_dim * 2, hidden_dim, kernel_size=4),
+            self.make_gen_block(hidden_dim, im_chan, kernel_size=4, final_layer=True),
+        )
+
+    def make_gen_block(self, input_channels, output_channels, kernel_size=3, stride=2, final_layer=False):
+        '''
+        Function to return a sequence of operations corresponding to a Decoder block of the VAE, 
+        corresponding to a transposed convolution, a batchnorm (except for in the last layer), and an activation
+        Parameters:
+        input_channels: how many channels the input feature representation has
+        output_channels: how many channels the output feature representation should have
+        kernel_size: the size of each convolutional filter, equivalent to (kernel_size, kernel_size)
+        stride: the stride of the convolution
+        final_layer: whether we're on the final layer (affects activation and batchnorm)
+        '''
+        if not final_layer:
+            return nn.Sequential(
+                nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
+                nn.BatchNorm2d(output_channels),
+                nn.ReLU(inplace=True),
+            )
+        else:
+            return nn.Sequential(
+                nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
+                nn.Sigmoid(),
+            )
+
+    def forward(self, noise):
+        '''
+        Function for completing a forward pass of the Decoder: Given a noise vector, 
+        returns a generated image.
+        Parameters:
+        noise: a noise tensor with dimensions (batch_size, z_dim)
+        '''
+        x = noise.view(len(noise), self.z_dim, 1, 1)
+        return self.gen(x)
+    
+
+class VAE(nn.Module):
+    '''
+    VAE Class
+    Values:
+    z_dim: the dimension of the noise vector, a scalar
+    im_chan: the number of channels of the output image, a scalar
+            MNIST is black-and-white, so that's our default
+    hidden_dim: the inner dimension, a scalar
+    '''
+    
+    def __init__(self, z_dim=Z_DIM, im_chan=3):
+        super(VAE, self).__init__()
+        self.z_dim = z_dim
+        self.encode = Encoder(im_chan, z_dim)
+        self.decode = Decoder(z_dim, im_chan)
+
+    def forward(self, images):
+        '''
+        Function for completing a forward pass of the Decoder: Given a noise vector, 
+        returns a generated image.
+        Parameters:
+        images: an image tensor with dimensions (batch_size, im_chan, im_height, im_width)
+        Returns:
+        decoding: the autoencoded image
+        q_dist: the z-distribution of the encoding
+        '''
+        q_mean, q_stddev = self.encode(images)
+        q_dist = Normal(q_mean, q_stddev)
+        z_sample = q_dist.rsample() # Sample once from each distribution, using the `rsample` notation
+        decoding = self.decode(z_sample)
+        return decoding, q_dist

embedding_models/vae/train.ipynb

@@ -0,0 +1,393 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import numpy as np\n",
+    "\n",
+    "from pathlib import Path\n",
+    "import os\n",
+    "from PIL import Image\n",
+    "\n",
+    "from model import VAE\n",
+    "from losses import *"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torch.utils.data import DataLoader, Dataset\n",
+    "from torchvision import transforms\n",
+    "import pandas as pd\n",
+    "import re\n",
+    "from sklearn.model_selection import train_test_split"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "IMAGE_FOLDER = './data/images/'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_names = os.listdir(IMAGE_FOLDER)\n",
+    "data = pd.DataFrame({'image_name': image_names})\n",
+    "data['label'] = data['image_name'].apply(lambda x: int(re.match('^\\d+', x)[0]))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class StampDataset(Dataset):\n",
+    "    def __init__(self, data, image_folder=Path(IMAGE_FOLDER), transform=None):\n",
+    "        super().__init__()\n",
+    "        self.image_folder = image_folder\n",
+    "        self.data = data\n",
+    "        self.transform = transform\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        image = Image.open(self.image_folder / self.data.iloc[idx]['image_name'])\n",
+    "        label = self.data.iloc[idx]['label']\n",
+    "        if self.transform:\n",
+    "            image = self.transform(image)\n",
+    "\n",
+    "        return image, label\n",
+    "\n",
+    "    \n",
+    "    def __len__(self):\n",
+    "        return len(self.data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_data, val_data = train_test_split(data, test_size=0.3, shuffle=True, stratify=data['label'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_transform = transforms.Compose([\n",
+    "    transforms.Resize((118, 118)),\n",
+    "    transforms.RandomHorizontalFlip(0.5),\n",
+    "    transforms.RandomVerticalFlip(0.5),\n",
+    "    transforms.ToTensor(),\n",
+    "    # transforms.Normalize((0.76302232, 0.77820438, 0.81879729), (0.16563211, 0.14949341, 0.1055889)),\n",
+    "])\n",
+    "\n",
+    "val_transform = transforms.Compose([\n",
+    "    transforms.Resize((118, 118)),\n",
+    "    transforms.ToTensor(),\n",
+    "    # transforms.Normalize((0.76302232, 0.77820438, 0.81879729), (0.16563211, 0.14949341, 0.1055889)),\n",
+    "])\n",
+    "train_dataset = StampDataset(train_data, transform=train_transform)\n",
+    "val_dataset = StampDataset(val_data, transform=val_transform)\n",
+    "\n",
+    "train_loader = DataLoader(train_dataset, shuffle=True, batch_size=256)\n",
+    "val_loader = DataLoader(val_dataset, shuffle=True, batch_size=256)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pytorch_lightning as pl\n",
+    "from torch import optim\n",
+    "from pytorch_lightning.loggers import TensorBoardLogger\n",
+    "\n",
+    "from torchvision.utils import make_grid"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "MEAN = torch.tensor((0.76302232, 0.77820438, 0.81879729)).view(3, 1, 1)\n",
+    "STD = torch.tensor((0.16563211, 0.14949341, 0.1055889)).view(3, 1, 1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class LitModel(pl.LightningModule):\n",
+    "    def __init__(self, alpha=1e-3):\n",
+    "        super().__init__()\n",
+    "        self.vae = VAE()\n",
+    "        self.vae_loss = VAELoss()\n",
+    "        self.triplet_loss = BatchHardTripletLoss(margin=1.)\n",
+    "        self.alpha = alpha\n",
+    "    \n",
+    "    def forward(self, x):\n",
+    "        return self.vae(x)\n",
+    "    \n",
+    "    def configure_optimizers(self):\n",
+    "        optimizer = optim.AdamW(self.parameters(), lr=3e-4)\n",
+    "        return optimizer\n",
+    "        # scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, 1000)\n",
+    "        # return {\"optimizer\": optimizer, \"lr_scheduler\": scheduler}\n",
+    "\n",
+    "    def training_step(self, batch, batch_idx):\n",
+    "        images, labels = batch\n",
+    "        labels = labels.unsqueeze(1)\n",
+    "        recon_images, encoding = self.vae(images)\n",
+    "        vae_loss = self.vae_loss(recon_images, images, encoding)\n",
+    "        self.log(\"train_vae_loss\", vae_loss, on_step=False, on_epoch=True, prog_bar=False, logger=True)\n",
+    "        triplet_loss = self.triplet_loss(encoding.mean, labels)\n",
+    "        self.log(\"train_triplet_loss\", triplet_loss, on_step=False, on_epoch=True, prog_bar=False, logger=True)\n",
+    "        loss = self.alpha * triplet_loss + vae_loss\n",
+    "        self.log(\"train_loss\", loss, on_epoch=True, prog_bar=True, logger=True)\n",
+    "        return loss\n",
+    "\n",
+    "    def validation_step(self, batch, batch_idx):\n",
+    "        images, labels = batch\n",
+    "        labels = labels.unsqueeze(1)\n",
+    "        recon_images, encoding = self.vae(images)\n",
+    "        vae_loss = self.vae_loss(recon_images, images, encoding)\n",
+    "        self.log(\"val_vae_loss\", vae_loss, on_step=False, on_epoch=True, prog_bar=False, logger=True)\n",
+    "        triplet_loss = self.triplet_loss(encoding.mean, labels)\n",
+    "        self.log(\"val_triplet_loss\", triplet_loss, on_step=False, on_epoch=True, prog_bar=False, logger=True)\n",
+    "        loss = self.alpha * triplet_loss + vae_loss\n",
+    "        self.log(\"val_loss\", loss, on_epoch=True, prog_bar=True, logger=True)\n",
+    "        return loss\n",
+    "\n",
+    "    def on_validation_epoch_end(self):\n",
+    "        images, _ = iter(val_loader).next()\n",
+    "        image_unflat = images.detach().cpu()\n",
+    "        image_grid = make_grid(image_unflat[:16], nrow=4)\n",
+    "        self.logger.experiment.add_image('original images', image_grid, self.current_epoch)\n",
+    "\n",
+    "        recon_images, _ = self.vae(images.to(self.device))\n",
+    "        image_unflat = recon_images.detach().cpu()\n",
+    "        image_grid = make_grid(image_unflat[:16], nrow=4)\n",
+    "        self.logger.experiment.add_image('reconstructed images', image_grid, self.current_epoch)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "litmodel = LitModel()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "logger = TensorBoardLogger(\"reconstruction_logs\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "epochs = 100"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = pl.Trainer(accelerator=\"auto\", max_epochs=epochs, logger=logger)\n",
+    "trainer.fit(model=litmodel, train_dataloaders=train_loader, val_dataloaders=val_loader)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%tensorboard"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import hf_hub_download"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "emb_model = torch.jit.load(hf_hub_download(repo_id=\"stamps-labs/vits8-stamp\", filename=\"vits8stamp-torchscript.pth\")).to(device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "val_transform = transforms.Compose([\n",
+    "    transforms.Resize((224, 224)),\n",
+    "    transforms.ToTensor(),\n",
+    "    # transforms.Normalize((0.76302232, 0.77820438, 0.81879729), (0.16563211, 0.14949341, 0.1055889)),\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_data['embed'] = train_data['image_name'].apply(lambda x: emb_model(val_transform(Image.open(Path(IMAGE_FOLDER) / x)).unsqueeze(0).to(device))[0].tolist())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\javid\\AppData\\Local\\Temp\\ipykernel_23064\\1572292890.py:1: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
+      "  embeds = pd.DataFrame(train_data['embed'].tolist()).append(pd.DataFrame(val_data['embed'].tolist()), ignore_index=True)\n",
+      "C:\\Users\\javid\\AppData\\Local\\Temp\\ipykernel_23064\\1572292890.py:2: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
+      "  labels = pd.DataFrame(train_data['label']).append(pd.DataFrame(val_data['label']), ignore_index=True)\n"
+     ]
+    }
+   ],
+   "source": [
+    "embeds = pd.DataFrame(train_data['embed'].tolist()).append(pd.DataFrame(val_data['embed'].tolist()), ignore_index=True)\n",
+    "labels = pd.DataFrame(train_data['label']).append(pd.DataFrame(val_data['label']), ignore_index=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "embeds.to_csv('./all_embeds.tsv', sep='\\t', index=False, header=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "labels.to_csv('./all_labels.tsv', sep='\\t', index=False, header=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 126,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "torch.save(litmodel.vae.encode.state_dict(), './models/encoder.pth')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 129,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "im = train_dataset[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 132,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<All keys matched successfully>"
+      ]
+     },
+     "execution_count": 132,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model = Encoder()\n",
+    "model.load_state_dict(torch.load('./models/encoder.pth'))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.9.0"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

embedding_models/vits8/example.py

@@ -0,0 +1,10 @@
+from PIL import Image
+from model import ViTStamp
+def get_embeddings(img_path: str):
+    model = ViTStamp()
+    image = Image.open(img_path)
+    embeddings = model(image=image)
+    return embeddings
+
+if __name__ == "__main__":
+    print(get_embeddings("oml/data/test/images/99d_15.bmp"))

embedding_models/vits8/model.py

@@ -0,0 +1,13 @@
+import torch
+from torchvision import transforms
+from huggingface_hub import hf_hub_download
+
+class ViTStamp():
+    def __init__(self):
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.model = torch.jit.load(hf_hub_download(repo_id="stamps-labs/vits8-stamp", filename="vits8stamp-torchscript.pth"))
+        self.transform = transforms.ToTensor()
+    def __call__(self, image) -> torch.Tensor():
+        img_tensor = self.transform(image).cuda().unsqueeze(0) if self.device == "cuda" else self.transform(image).unsqueeze(0)
+        features = self.model(img_tensor)
+        return features

embedding_models/vits8/oml/create_dataset.py

@@ -0,0 +1,71 @@
+import os
+from PIL import Image
+import pandas as pd
+
+import argparse
+
+parser = argparse.ArgumentParser("Create a dataset for training with OML", 
+                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+parser.add_argument("--root-data-path", help="Path to images for dataset", default="data/train_val/")
+parser.add_argument("--image-data-path", help="Image folder in root data path", default="images/")
+parser.add_argument("--train-val-split", 
+                    help="In which ratio to split data in format train:val (For example 80:20)", default="80:20")
+parser.add_argument("--separator", 
+                    help="What separator is used in image name to separate class name and instance (E.g. circle1_5, separator=_)",
+                    default="_")
+
+args = parser.parse_args()
+config = vars(args)
+
+root_path = config["root_data_path"]
+image_path = config["image_data_path"]
+separator = config["separator"]
+
+train_prc, val_prc = tuple(int(num)/100 for num in config["train_val_split"].split(":"))
+
+class_names = set()
+for image in os.listdir(root_path+image_path):
+    if image.endswith(("png", "jpg", "bmp", "webp")):
+        img_name = image.split(".")[0]
+        Image.open(root_path+image_path+image).resize((224,224)).save(root_path+image_path+img_name+".png", "PNG")
+        if not image.endswith("png"):
+            os.remove(root_path+image_path+image)
+        img_name = img_name.split(separator)
+        class_name = img_name[0]+img_name[1]
+        class_names.add(class_name)
+    else:
+        print("Not all of the images are in supported format")
+    
+
+#For each class in set assign its index in a set as a class label.
+class_label_dict = {}
+for ind, name in enumerate(class_names):
+    class_label_dict[name] = ind
+
+class_count = len(class_names)
+train_class_count = int(class_count*train_prc)
+print(train_class_count)
+
+df_dict = {"label": [],
+           "path": [],
+           "split": [],
+           "is_query": [],
+           "is_gallery": []}
+for image in os.listdir(root_path+image_path):
+    if image.endswith((".png", ".jpg", ".bmp", ".webp")):
+        img_name = image.split(".")[0].split(separator)
+        class_name = img_name[0]+img_name[1]
+        label = class_label_dict[class_name]
+        path = image_path+image
+        split = "train" if label <= train_class_count else "validation"
+        is_query, is_gallery = (1, 1) if split=="validation" else (None, None)
+        df_dict["label"].append(label)
+        df_dict["path"].append(path)
+        df_dict["split"].append(split)
+        df_dict["is_query"].append(is_query)
+        df_dict["is_gallery"].append(is_gallery)
+
+df = pd.DataFrame(df_dict)
+
+df.to_csv(root_path+"df_stamps.csv", index=False)

embedding_models/vits8/oml/data/train_val/df_stamps.csv

@@ -0,0 +1,41 @@
+label,path,split,is_query,is_gallery
+0,images/circle6_1239.png,train,,
+8,images/triangle19_1242.png,train,,
+21,images/rectangle11_1248.png,train,,
+39,images/triangle10_1232.png,validation,1.0,1.0
+33,images/word14_1241.png,validation,1.0,1.0
+38,images/word5_1233.png,validation,1.0,1.0
+15,images/circle19_1236.png,train,,
+22,images/circle15_1244.png,train,,
+32,images/circle21_1249.png,train,,
+26,images/oval20_1242.png,train,,
+6,images/oval5_1237.png,train,,
+23,images/word9_1241.png,train,,
+9,images/triangle22_1238.png,train,,
+31,images/circle12_1239.png,train,,
+11,images/word21_1231.png,train,,
+4,images/oval2_1235.png,train,,
+20,images/rectangle18_1246.png,train,,
+12,images/circle24_1234.png,train,,
+5,images/circle2_1249.png,train,,
+37,images/word22_1238.png,validation,1.0,1.0
+34,images/triangle18_1247.png,validation,1.0,1.0
+1,images/oval7_1241.png,train,,
+10,images/triangle13_1240.png,train,,
+14,images/rectangle12_1236.png,train,,
+36,images/circle8_1237.png,validation,1.0,1.0
+24,images/triangle9_1245.png,train,,
+29,images/word23_1243.png,train,,
+28,images/triangle11_1244.png,train,,
+16,images/circle2_1246.png,train,,
+30,images/circle3_1247.png,train,,
+18,images/oval24_1248.png,train,,
+2,images/oval12_1231.png,train,,
+3,images/oval18_1234.png,train,,
+25,images/rectangle11_1245.png,train,,
+17,images/word9_1244.png,train,,
+13,images/triangle14_1237.png,train,,
+35,images/circle2_1233.png,validation,1.0,1.0
+7,images/word18_1239.png,train,,
+19,images/rectangle13_1236.png,train,,
+27,images/circle24_1246.png,train,,