Spaces:
Runtime error
Runtime error
File size: 7,180 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 164 165 166 167 168 169 170 171 |
"""
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 tqdm import tqdm
def main(params):
if params["model_type"] != 'vit_base_patch32_224_in21k':
model, transform = load(params["model_type"], jit=False)
else:
model = timm.create_model(params["model_type"], pretrained=True)
model = model.cuda()
if params["model_type"] != 'vit_base_patch32_224_in21k':
save_model_type = params["model_type"].split("-")[0]
mean = torch.Tensor([0.48145466, 0.4578275, 0.40821073]).to("cuda").reshape(3, 1, 1)
std = torch.Tensor([0.26862954, 0.26130258, 0.27577711]).to("cuda").reshape(3, 1, 1)
if "RN" in params["model_type"]:
num_patches = 196 #600 * 1000 // 32 // 32
pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, model.visual.attnpool.positional_embedding.shape[-1], device='cuda'),)
pos_embed.weight = resize_pos_embed(model.visual.attnpool.positional_embedding.unsqueeze(0), pos_embed)
model.visual.attnpool.positional_embedding = pos_embed
else:
save_model_type = 'vit_base'
mean = torch.Tensor([0.5, 0.5, 0.5]).to("cuda").reshape(3, 1, 1)
std = torch.Tensor([0.5, 0.5, 0.5]).to("cuda").reshape(3, 1, 1)
num_patches = 196 #600 * 1000 // 32 // 32
pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, 768, device='cuda'),)
pos_embed.weight = resize_pos_embed(model.pos_embed, pos_embed)
model.pos_embed = pos_embed
if params["model_type"] == "ViT-B/32":
num_patches = 196 #600 * 1000 // 32 // 32
pos_embed = nn.Parameter(torch.zeros(num_patches + 1, 768, device='cuda'),)
pos_embed.weight = resize_pos_embed(model.visual.positional_embedding.unsqueeze(0), pos_embed.unsqueeze(0))
model.visual.positional_embedding = pos_embed
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'
if not os.path.isdir(dir_fc):
os.mkdir(dir_fc)
if not os.path.isdir(dir_att):
os.mkdir(dir_att)
for i,img in enumerate(tqdm(imgs)):
# load the image
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(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 ', dir_fc, dir_att)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# input 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')
parser.add_argument('--batch_size', default=1, type=int, help='batch size')
args = parser.parse_args()
params = vars(args) # convert to ordinary dict
print('parsed input parameters:')
print(json.dumps(params, indent = 2))
main(params)
|