KalharaX/VPR

1import os2import math3import numpy as np4import torch as th5import torch.nn as nn6import torch.nn.functional as F7import open_clip8from torch.utils.data import Dataset9from torch.utils.data import DataLoader10import torchvision.transforms as T11from tqdm import tqdm12import pandas as pd13import cv214from PIL import Image15from get_image import *16 17 18path1 = 'model_epoch_4_mAP3_0.34.pt'19 20def read_img(img_path, is_gray=False):21    mode = cv2.IMREAD_COLOR if not is_gray else cv2.IMREAD_GRAYSCALE22    img = cv2.imread(img_path, mode)23    if not is_gray:24        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)25    return img26 27class ProductDataset(Dataset):28    def __init__(self,29                 img_dir,30                 annotations_file,31                 transform=None,32                 final_transform=None,33                 headers=None,34                 test_mode=False):35        self.data = pd.read_csv(annotations_file)36        self.img_dir = img_dir37        self.transform = transform38        self.final_transform = final_transform39        self.headers = {"img_path": "img_path", "product_id": "product_id"}40        if headers:41            self.headers = headers42        self.test_mode = test_mode43 44    def __len__(self):45        return len(self.data)46 47    def __getitem__(self, idx):48        img_path = os.path.join(self.img_dir, self.data[self.headers["img_path"]][idx])49 50        img = read_img(img_path)51        if self.test_mode:52            x, y, w, h = self.data["bbox_x"][idx], self.data["bbox_y"][idx], \53                         self.data["bbox_w"][idx], self.data["bbox_h"][idx]54            img = img[y:y+h, x:x+w]55 56 57        if self.transform is not None:58            img = transform(image=img)["image"]59 60        if self.final_transform is not None:61            if isinstance(img, np.ndarray):62                img =  Image.fromarray(img)63            img = self.final_transform(img)64 65        product_id = self.data[self.headers["product_id"]][idx]66        return img, product_id67 68def get_final_transform():69    final_transform = T.Compose([70            T.Resize(71                size=(224, 224),72                interpolation=T.InterpolationMode.BICUBIC,73                antialias=True),74            T.ToTensor(),75            T.Normalize(76                mean=(0.48145466, 0.4578275, 0.40821073),77                std=(0.26862954, 0.26130258, 0.27577711)78            )79        ])80    return final_transform81 82@th.no_grad()83def extract_embeddings(model, dataloader, epoch=10, use_cpu=True):84    features = []85    product_id = []86 87    for _ in range(epoch):88        for imgs, p_id in tqdm(dataloader):89            if use_cpu:90                imgs = imgs.cuda()91            features.append(th.squeeze(model(imgs.half())).detach().cpu().numpy().astype(np.float32))92            product_id.append(th.squeeze(p_id).detach().cpu().numpy())93 94 95    return np.concatenate(features, axis=0), np.concatenate(product_id)96 97 98 99 100 101 102 103 104 105