Spaces:
Build error
Build error
# make.texts.py | |
from __future__ import print_function | |
import os | |
import os.path as osp | |
from pycocotools.coco import COCO | |
# import gensim | |
# from gensim.models import Doc2Vec | |
import numpy as np | |
import scipy.io as sio | |
import os | |
import os.path as osp | |
from pycocotools.coco import COCO | |
import pprint | |
import os | |
import os.path as osp | |
import json | |
from nltk.tokenize import RegexpTokenizer | |
from tqdm import tqdm | |
"""process texts | |
python 2 needed by `jhlau/doc2vec`, and COCO api CAN work with python 2.7. | |
So I choose to create a virtual env of python 2.7. | |
dependencies: | |
matplotlib (COCO api) | |
smart_open (gensim) | |
""" | |
# COCO 原本的 annotations 中就有各 classes 的 ID,但不连续(从 1 标到 90 但实际只有 80 个)。这里按原有的 category id 的升序重新定义连续的、0-based 的 class ID。 | |
# train 和 val 都包含所有类,所以这里只用 val set 处理。 | |
# 结果写入 class-name.COCO.txt | |
def remake_classname(): | |
"""process class order | |
Record the mapping between tightened/discretized 0-base class ID, | |
original class ID and class name in `class-name.COCO.txt`, | |
with format `<new ID> <original ID> <class name>`. | |
The class order is consistent to the ascending order of the original IDs. | |
""" | |
COCO_P = "/dataset/coco" | |
ANNO_P = osp.join(COCO_P, "annotations") | |
SPLIT = ["val", "train"] | |
for _split in SPLIT: | |
print("---", _split, "---") | |
anno_file = osp.join(ANNO_P, "instances_{}2017.json".format(_split)) | |
coco = COCO(anno_file) | |
cats = coco.loadCats(coco.getCatIds()) | |
# print(cats[0]) | |
cls_id = {c["name"]: c["id"] for c in cats} # 它本身就是按 category id 升序 | |
# pprint.pprint(cls_id) | |
with open("class-name.COCO.txt", "w") as f: | |
for new_id, c in enumerate(cls_id): | |
old_id = cls_id[c]# - 1 | |
cn = c.replace(" ", "_") | |
# format: <new ID> <original ID> <class name> | |
f.write("{} {} {}\n".format(new_id, old_id, cn)) | |
break # 只用 val set | |
def remake_idmap(): | |
# 合并 train、val 两个集合,统一按原本的 id(即 images 文件名中的数字,也是不连续的,且 train、val 无重合)升序重新排 0-based 的 data ID。 | |
# 结果写入 id-map.COCO.txt | |
# make.id-map.py | |
"""discretization of the original file ID | |
Map the file ID to sequential {0, 1, ..., n}, | |
and record this mapping in `id-map.txt`, | |
with format `<new id> <original id> <image file name>`. | |
Note that the new ids are 0-base. | |
""" | |
TRAIN_P = "train2017" | |
VAL_P = "val2017" | |
file_list = [f for f in os.listdir(os.path.join("/dataset/coco", TRAIN_P)) if (".jpg" in f)] | |
file_list.extend([f for f in os.listdir(os.path.join("/dataset/coco", VAL_P)) if (".jpg" in f)]) | |
print("#data:", len(file_list)) # 12,3287 | |
id_key = lambda x: int(x.split(".jpg")[0]) | |
file_list = sorted(file_list, key=id_key) # 按 image ID 升序 | |
# print(file_list[:15]) | |
with open("id-map.COCO.txt", "w") as f: | |
# format: <new id> <original id> <image file name> | |
for i, f_name in enumerate(file_list): | |
_original_id = id_key(f_name) | |
f.write("{} {} {}\n".format(i, _original_id, f_name)) | |
# if i > 5: break | |
print("DONE") | |
# COCO | |
COCO_P = "/dataset/coco" | |
ANNO_P = osp.join(COCO_P, "annotations") | |
SPLIT = ["val", "train"] | |
# doc2vec | |
MODEL = "/home/dataset/Doc2Vec/enwiki_dbow/doc2vec.bin" | |
start_alpha = 0.01 | |
infer_epoch = 1000 | |
DIM = 300 # dimension of the doc2vec feature | |
# id_map_data = {} | |
# with open("id-map.txt", "r") as f: | |
# for line in f: | |
# line = line.strip() | |
# _new_id, _old_id, _ = line.split() | |
# id_map_data[int(_old_id)] = int(_new_id) | |
# N_DATA = len(id_map_data) | |
# print("#data:", N_DATA) | |
# pre-trained Doc2Vec model | |
# model = Doc2Vec.load(MODEL) | |
tokenizer = RegexpTokenizer(r'\w+') | |
def dataset_format(filepath, filename, imgid, split, sentences, cocoid): | |
data = {} | |
data['filepath'] = filepath | |
data['sentids'] = [imgid * 5 + idx for idx in range(5)] | |
data['filename'] = filename | |
data['imgid'] = imgid | |
data['split'] = split | |
data['sentences'] = [{'tokens': tokenizer.tokenize(sentence), | |
'raw': sentence, | |
'imgid': imgid, | |
'sentid': imgid * 5 + idx} | |
for idx, sentence in enumerate(sentences)] | |
data['cocoid'] = cocoid | |
return data | |
dataset_anns = {} | |
dataset_anns['images'] = [] | |
dataset_anns['dataset'] = 'coco' | |
for __split in SPLIT: | |
print("---", __split, "---") | |
anno_file = osp.join(ANNO_P, "instances_{}2017.json".format(__split)) | |
caps_file = osp.join(ANNO_P, "captions_{}2017.json".format(__split)) | |
coco = COCO(anno_file) | |
coco_caps = COCO(caps_file) | |
new_image_id_file = open("id-map.COCO.txt", 'r') | |
new_img_id_map = {image_id.strip().split(" ")[2]: image_id.strip().split(" ")[0] for image_id in new_image_id_file.readlines()} | |
id_list = coco.getImgIds() | |
for _old_id in tqdm(id_list): | |
# _new_id = id_map_data[_old_id] | |
_annIds = coco_caps.getAnnIds(imgIds=_old_id) | |
_anns = coco_caps.loadAnns(_annIds) | |
_filepath = __split + '2017' | |
_filename = coco.imgs[_old_id]['file_name'] | |
_imgid = int(new_img_id_map[_filename]) | |
_split = __split | |
# print(len(anns)) | |
# pprint.pprint(anns) | |
_sentences = [_a["caption"] for _a in _anns] | |
_cocoid = _old_id | |
formated_data = dataset_format(_filepath, _filename, _imgid, _split, _sentences, _cocoid) | |
dataset_anns['images'].append(formated_data) | |
# pprint.pprint(sentences) | |
# sentences = [gensim.utils.simple_preprocess(s) for s in sentences] | |
# pprint.pprint(sentences) | |
# doc = [] | |
# for s in sentences: | |
# doc.extend(s) | |
# print(doc) | |
# vec = model.infer_vector(doc) | |
# print(vec.shape) | |
# texts.append(vec[np.newaxis, :]) | |
# break | |
# break | |
with open('dataset_anns.json', 'w') as fp: | |
json.dump(dataset_anns, fp) | |
new_image_id_file.close() | |
# texts = np.vstack(texts).astype(np.float32) | |
# print("texts:", texts.shape, texts.dtype) # (123287, 300) dtype('<f4') | |
# sio.savemat("texts.COCO.doc2vec.{}.mat".format(DIM), {"texts": texts}) | |