Spaces:
Runtime error
Runtime error
File size: 6,183 Bytes
c80917c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
"""
Preprocess a raw json dataset into features files for use in data_loader.py
Input: json file that has the form
[{ file_path: 'path/img.jpg', captions: ['a caption', ...] }, ...]
example element in this list would look like
{'captions': [u'A man with a red helmet on a small moped on a dirt road. ', u'Man riding a motor bike on a dirt road on the countryside.', u'A man riding on the back of a motorcycle.', u'A dirt path with a young person on a motor bike rests to the foreground of a verdant area with a bridge and a background of cloud-wreathed mountains. ', u'A man in a red shirt and a red hat is on a motorcycle on a hill side.'], 'file_path': u'val2014/COCO_val2014_000000391895.jpg', 'id': 391895}
This script reads this json, does some basic preprocessing on the captions
(e.g. lowercase, etc.), creates a special UNK token, and encodes everything to arrays
Output: two folders of features
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import json
import argparse
from random import shuffle, seed
import string
# non-standard dependencies:
import h5py
from six.moves import cPickle
import numpy as np
import torch
import torchvision.models as models
import skimage.io
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
from PIL import Image
from torch import nn
# preprocess = Compose([
# Resize((448, 448), interpolation=Image.BICUBIC),
# CenterCrop((448, 448)),
# ToTensor()
# ])
# from clip.clip import load
# from timm.models.vision_transformer import resize_pos_embed
# import timm
# from captioning.utils.resnet_utils import myResnet
# import captioning.utils.resnet as resnet
from captioning.utils.clipscore import CLIPScore
from tqdm import tqdm
def main(params):
clipscore_model = CLIPScore()
clipscore_model.to('cuda')
imgs = json.load(open(params['input_json'], 'r'))
imgs = imgs['images']
if args.n_jobs > 1:
print('Total imgs:', len(imgs))
print('Using {} jobs'.format(args.n_jobs))
print('job id:', args.job_id)
imgs = imgs[args.job_id::args.n_jobs]
N = len(imgs)
seed(123) # make reproducible
# dir_fc = params['output_dir']+'_clip_'+save_model_type+'_fc'
# dir_att = params['output_dir']+'_clip_'+save_model_type+'_att'
vis_dir_fc = params['output_dir']+'_clipscore_vis'
if not os.path.isdir(vis_dir_fc):
os.mkdir(vis_dir_fc)
# text_dir_fc = params['output_dir']+'_clipscore_text'
# if not os.path.isdir(text_dir_fc):
# os.mkdir(text_dir_fc)
# if not os.path.isdir(dir_att):
# os.mkdir(dir_att)
for i, img in enumerate(tqdm(imgs)):
# load the image
img_path = os.path.join(params['images_root'], img['filepath'], img['filename'])
img_feat = clipscore_model.image_extract(img_path)
img_feat = img_feat.view(512)
# for d in img['sentences']:
# text = d['raw'].strip()
# text_feat = clipscore_model.text_extract(text)
# with torch.no_grad():
# image = preprocess(Image.open(os.path.join(params['images_root'], img['filepath'], img['filename']) ).convert("RGB"))
# image = torch.tensor(np.stack([image])).cuda()
# image -= mean
# image /= std
# if "RN" in params["model_type"]:
# tmp_att, tmp_fc = model.encode_image(image)
# tmp_att = tmp_att[0].permute(1, 2, 0)
# tmp_fc = tmp_fc[0]
# elif params["model_type"] == 'vit_base_patch32_224_in21k':
# x = model(image)
# tmp_fc = x[0, 0, :]
# tmp_att = x[0, 1:, :].reshape( 14, 14, 768 )
# else:
# x = model.visual.conv1(image.half()) # shape = [*, width, grid, grid]
# x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2]
# x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width]
# x = torch.cat([model.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1) # shape = [*, grid ** 2 + 1, width]
# x = x + model.visual.positional_embedding.to(x.dtype)[:x.shape[1], :]
# x = model.visual.ln_pre(x)
# x = x.permute(1, 0, 2) # NLD -> LND
# for layer_idx, layer in enumerate(model.visual.transformer.resblocks):
# x = layer(x)
# x = x.permute(1, 0, 2)
# tmp_fc = x[0, 0, :]
# tmp_att = x[0, 1:, :].reshape( 14, 14, 768 )
np.save(os.path.join(vis_dir_fc, str(img['cocoid'])), img_feat.data.cpu().float().numpy())
# np.save(os.path.join(text_dir_fc, str(img['cocoid'])), tmp_fc.data.cpu().float().numpy())
# np.savez_compressed(os.path.join(dir_att, str(img['cocoid'])), feat=tmp_att.data.cpu().float().numpy())
if i % 1000 == 0:
print('processing %d/%d (%.2f%% done)' % (i, N, i*100.0/N))
print('wrote ', vis_dir_fc)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# input json
# dataset_coco.json
parser.add_argument('--input_json', required=True, help='input json file to process into hdf5')
parser.add_argument('--output_dir', default='data', help='output h5 file')
# options
parser.add_argument('--images_root', default='', help='root location in which images are stored, to be prepended to file_path in input json')
# parser.add_argument('--att_size', default=14, type=int, help='14x14 or 7x7')
# parser.add_argument('--model_type', default='RN50', type=str, help='RN50, RN101, RN50x4, ViT-B/32, vit_base_patch32_224_in21k')
parser.add_argument('--n_jobs', default=-1, type=int, help='number of jobs to run in parallel')
parser.add_argument('--job_id', default=0, type=int, help='job id')
args = parser.parse_args()
params = vars(args) # convert to ordinary dict
print('parsed input parameters:')
print(json.dumps(params, indent = 2))
main(params)
|