Spaces:
Runtime error
Runtime error
| """ | |
| # 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 | |
| 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) | |
| # preprocessors | |
| vis_processor = BlipImageEvalProcessor(image_size=224) | |
| text_processor = BlipCaptionProcessor() | |
| # files to process | |
| # file_root = "/export/home/.cache/lavis/coco/images/val2014" | |
| 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") | |