|
"""
|
|
# Copyright (c) 2022, salesforce.com, inc.
|
|
# All rights reserved.
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
# For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
|
|
"""
|
|
|
|
from PIL import Image
|
|
import requests
|
|
import torch
|
|
|
|
import os
|
|
|
|
from lavis.common.registry import registry
|
|
from lavis.processors import *
|
|
from lavis.models import *
|
|
from lavis.common.utils import build_default_model
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
def load_demo_image():
|
|
img_url = (
|
|
"https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg"
|
|
)
|
|
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
|
|
|
|
return raw_image
|
|
|
|
|
|
def read_img(filepath):
|
|
raw_image = Image.open(filepath).convert("RGB")
|
|
|
|
return raw_image
|
|
|
|
|
|
|
|
model_url = "https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base.pth"
|
|
feature_extractor = BlipFeatureExtractor(pretrained=model_url)
|
|
|
|
feature_extractor.eval()
|
|
feature_extractor = feature_extractor.to(device)
|
|
|
|
|
|
vis_processor = BlipImageEvalProcessor(image_size=224)
|
|
text_processor = BlipCaptionProcessor()
|
|
|
|
|
|
|
|
file_root = "/export/home/.cache/lavis/coco/images/train2014"
|
|
filepaths = os.listdir(file_root)
|
|
|
|
print(len(filepaths))
|
|
|
|
caption = "dummy"
|
|
|
|
path2feat = dict()
|
|
bsz = 256
|
|
|
|
images_in_batch = []
|
|
filepaths_in_batch = []
|
|
|
|
for i, filename in enumerate(filepaths):
|
|
if i % bsz == 0 and i > 0:
|
|
images_in_batch = torch.cat(images_in_batch, dim=0).to(device)
|
|
with torch.no_grad():
|
|
image_features = feature_extractor(
|
|
images_in_batch, caption, mode="image", normalized=True
|
|
)[:, 0]
|
|
|
|
for filepath, image_feat in zip(filepaths_in_batch, image_features):
|
|
path2feat[os.path.basename(filepath)] = image_feat.detach().cpu()
|
|
|
|
images_in_batch = []
|
|
filepaths_in_batch = []
|
|
|
|
print(len(path2feat), image_features.shape)
|
|
else:
|
|
filepath = os.path.join(file_root, filename)
|
|
|
|
image = read_img(filepath)
|
|
image = vis_processor(image).unsqueeze(0)
|
|
|
|
images_in_batch.append(image)
|
|
filepaths_in_batch.append(filepath)
|
|
|
|
torch.save(path2feat, "path2feat_coco_train2014.pth")
|
|
|
