File size: 5,958 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
"""
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']
    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_path = os.path.join(params['images_root'], img['file_name'])
        img_path = os.path.join(params['images_root'], img['file_path'])

        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['id'])), 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')

    args = parser.parse_args()
    params = vars(args) # convert to ordinary dict
    print('parsed input parameters:')
    print(json.dumps(params, indent = 2))
    main(params)