|
import json |
|
from tqdm import tqdm |
|
import re |
|
import fire |
|
|
|
|
|
def tokenize_caption(input_json: str, |
|
keep_punctuation: bool = False, |
|
host_address: str = None, |
|
character_level: bool = False, |
|
zh: bool = True, |
|
output_json: str = None): |
|
"""Build vocabulary from csv file with a given threshold to drop all counts < threshold |
|
|
|
Args: |
|
input_json(string): Preprossessed json file. Structure like this: |
|
{ |
|
'audios': [ |
|
{ |
|
'audio_id': 'xxx', |
|
'captions': [ |
|
{ |
|
'caption': 'xxx', |
|
'cap_id': 'xxx' |
|
} |
|
] |
|
}, |
|
... |
|
] |
|
} |
|
threshold (int): Threshold to drop all words with counts < threshold |
|
keep_punctuation (bool): Includes or excludes punctuation. |
|
|
|
Returns: |
|
vocab (Vocab): Object with the processed vocabulary |
|
""" |
|
data = json.load(open(input_json, "r"))["audios"] |
|
|
|
if zh: |
|
from nltk.parse.corenlp import CoreNLPParser |
|
from zhon.hanzi import punctuation |
|
parser = CoreNLPParser(host_address) |
|
for audio_idx in tqdm(range(len(data)), leave=False, ascii=True): |
|
for cap_idx in range(len(data[audio_idx]["captions"])): |
|
caption = data[audio_idx]["captions"][cap_idx]["caption"] |
|
|
|
if not keep_punctuation: |
|
caption = re.sub("[{}]".format(punctuation), "", caption) |
|
if character_level: |
|
tokens = list(caption) |
|
else: |
|
tokens = list(parser.tokenize(caption)) |
|
data[audio_idx]["captions"][cap_idx]["tokens"] = " ".join(tokens) |
|
else: |
|
from pycocoevalcap.tokenizer.ptbtokenizer import PTBTokenizer |
|
captions = {} |
|
for audio_idx in range(len(data)): |
|
audio_id = data[audio_idx]["audio_id"] |
|
captions[audio_id] = [] |
|
for cap_idx in range(len(data[audio_idx]["captions"])): |
|
caption = data[audio_idx]["captions"][cap_idx]["caption"] |
|
captions[audio_id].append({ |
|
"audio_id": audio_id, |
|
"id": cap_idx, |
|
"caption": caption |
|
}) |
|
tokenizer = PTBTokenizer() |
|
captions = tokenizer.tokenize(captions) |
|
for audio_idx in tqdm(range(len(data)), leave=False, ascii=True): |
|
audio_id = data[audio_idx]["audio_id"] |
|
for cap_idx in range(len(data[audio_idx]["captions"])): |
|
tokens = captions[audio_id][cap_idx] |
|
data[audio_idx]["captions"][cap_idx]["tokens"] = tokens |
|
|
|
if output_json: |
|
json.dump( |
|
{ "audios": data }, open(output_json, "w"), |
|
indent=4, ensure_ascii=not zh) |
|
else: |
|
json.dump( |
|
{ "audios": data }, open(input_json, "w"), |
|
indent=4, ensure_ascii=not zh) |
|
|
|
|
|
if __name__ == "__main__": |
|
fire.Fire(tokenize_caption) |
|
|
