SPRM/make_dataset_obj_dict_caffe.py at main · drigoni/SPRM · GitHub

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
from __future__ import print_function
from collections import defaultdict
from email.mime import image

import os
from os import listdir
from os import path
from os.path import isfile, join
import argparse
from posixpath import split
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import base64
import csv
import h5py
import _pickle as cPickle
import numpy as np
import utils
import csv
import pandas as pd
import base64
import json


def save_json(data, output_folder, split_name):
    file_name = '{}_detection_dict.json'.format(split_name)
    file_path = os.path.join(output_folder, file_name)
    with open(file_path, 'w') as f:
        json.dump(data, f)
    print('Data saved: ', file_path)


def load_boxes_classes(file_classes, do_spellchecker=False):
	# read labels
	with open(file_classes, 'r') as f:
		data = f.readlines()
	labels = ["__background__"] + [label.strip() for label in data] # add background class

	# correct phrase
	if do_spellchecker:
		new_labels = []
		spell = SpellChecker()
		for label in labels:
			label_corrected = spell.correction(label)
			if label_corrected is not None:
				new_labels.append(label_corrected)
			else:
				new_labels.append(label)
	else:
		new_labels = labels
	return new_labels


def load_data(img_folder):
    class_labels = load_boxes_classes('data/objects_vocab.txt')
    attributes_labels = load_boxes_classes('data/attributes_vocab.txt')
    # get all file in the folder
    onlyfiles = [join(img_folder, f) for f in listdir(img_folder) if isfile(join(img_folder, f))]
    print('Number of files: ', len(onlyfiles))
    onlyfiles = [f for f in onlyfiles if f[-4:] == '.npz']
    print('Number of .npz files: ', len(onlyfiles))

    # load all data ['image_id', 'image_w', 'image_h', 'num_boxes', 'boxes', 'features']
    all_data = dict()
    for img_file in onlyfiles:
        img_id = img_file.split('/')[-1][:-8]   # remove ".jpg.npz"
        tmp_dict = dict()
        with np.load(img_file, allow_pickle=True) as f:
            # NOTE: background included in the extarcted classes. So we filter those boxes
            best_class_idx = np.argmax(f['cls_prob'], axis=-1)
            best_attribute_idx = np.argmax(f['attr_prob'], axis=-1)
            # filter
            selected_boxes = f['bbox'][best_class_idx > 0, :]
            selected_classes = best_class_idx[best_class_idx > 0]
            selected_attributes = best_attribute_idx[best_class_idx > 0]
            # retrieve labels
            best_class_labels = [class_labels[i] for i in selected_classes]
            best_attribute_labels = [attributes_labels[i] for i in selected_attributes]
            # save
            tmp_dict['bboxes'] = selected_boxes.tolist()
            tmp_dict['classes'] = best_class_labels
            tmp_dict['attrs'] = best_attribute_labels
        all_data[img_id] = tmp_dict
    return all_data


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--extracted_features', type=str, default='./data/referit/extracted_features/', help='Folder of extracted features.')
    parser.add_argument('--output_folder', type=str, default='./data/referit/', help='Folder where to save the detection_dict.json file.')
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = parse_args()
    if os.path.exists(args.extracted_features):
        print('Loading all data.')
        all_data = load_data(args.extracted_features)

        print("Saving data.")
        splits = ['./data/referit/train.txt',
                    './data/referit/val.txt',
                    './data/referit/test.txt']
        for split_path in splits:
            split_name = split_path.split('/')[-1][:-4]
            with open(split_path, 'r') as f:
                split_idx = f.read().splitlines()
            split_idx = [i.strip() for i in split_idx]  # cleaning
            print('Processing split {} with {} elements. '.format(split_name, len(all_data)))
            save_json(all_data, args.output_folder, split_name)
    else:
        print("Folder not valid: ", args.extracted_features)
        exit(1)