|
|
import glob |
|
|
import os |
|
|
import PIL |
|
|
import PIL.Image |
|
|
import numpy as np |
|
|
import torch |
|
|
from torch.utils.data import DataLoader |
|
|
from tqdm import tqdm |
|
|
import open_clip |
|
|
|
|
|
|
|
|
class SingleFolderDataset(torch.utils.data.Dataset): |
|
|
def __init__(self, folder_path, transform=None): |
|
|
self.folder_path = folder_path |
|
|
self.transform = transform |
|
|
self.image_paths = glob.glob(os.path.join(folder_path, "*")) |
|
|
print('Found {} images in {}'.format(len(self.image_paths), folder_path)) |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.image_paths) |
|
|
|
|
|
def __getitem__(self, index): |
|
|
image_path = self.image_paths[index] |
|
|
image = PIL.Image.open(image_path) |
|
|
if self.transform: |
|
|
image = self.transform(image) |
|
|
return image, os.path.basename(image_path) |
|
|
|
|
|
|
|
|
def extract_feats(index_config): |
|
|
|
|
|
ai_config = index_config['a1_config'] |
|
|
weight_path = index_config['weight_path'] |
|
|
img_dir = index_config['img_dir'] |
|
|
batch_size = 1024 |
|
|
|
|
|
model, _, transform = open_clip.create_model_and_transforms(ai_config, pretrained=weight_path) |
|
|
|
|
|
devive = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') |
|
|
print('Using device:', devive) |
|
|
model = model.to(devive) |
|
|
|
|
|
dataset = SingleFolderDataset(img_dir, transform=transform) |
|
|
|
|
|
dl = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4) |
|
|
|
|
|
im_ids = [] |
|
|
im_feats = [] |
|
|
|
|
|
for i, (patched_tensor, img_id) in tqdm(enumerate(dl)): |
|
|
patched_tensor = patched_tensor.to(devive) |
|
|
|
|
|
with torch.no_grad(): |
|
|
out = model.encode_image(patched_tensor) |
|
|
|
|
|
im_ids.append(img_id) |
|
|
im_feats.append(out.cpu().numpy()) |
|
|
|
|
|
im_hashes = np.concatenate(im_ids) |
|
|
im_feats = np.concatenate(im_feats) |
|
|
return im_hashes, im_feats |
|
|
|
|
|
|
|
|
|
