Serial-Speakers/encrypt_text.py at master · bostxavier/Serial-Speakers · 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
#!/usr/bin/python3
# -*- coding: utf-8 -*-

import sys
import argparse
import os
import json
import hashlib

from nltk.tokenize import word_tokenize


def encrypt_text(annot_file, output_file):
    """Encrypts the textual content of the ``Serial Speakers'' dataset
    Args:
    annot_file  (str): input annotation file with plain text
    output_file (str): output annotation file with encrypted text

    Returns:
    None
    """
    # expand input paths
    annot_file = os.path.expanduser(annot_file)
    output_file = os.path.expanduser(output_file)

    # load annotation file
    annotations = json.load(open(annot_file))

    # dictionary of hashes
    hash_dict = {}

    # loop over seasons
    seasons = annotations['seasons']
    for i, season in enumerate(seasons):

        # loop over episodes
        episodes = season['episodes']
        for j, episode in enumerate(episodes):

            annotations['seasons'][i]['episodes'][j].pop('path')
            annotations['seasons'][i]['episodes'][j].pop('width')
            annotations['seasons'][i]['episodes'][j].pop('height')

            # loop over speech segments
            speech_segments = episode['data']['speech_segments']

            # new list of encrypted speech segments
            new_speech_segments = []

            for speech_segment in speech_segments:
                text = speech_segment['text'].lower()

                # encrypt words
                encrypted_tokens = []
                words = word_tokenize(text)
                for word in words:

                    # new word type: compute hash
                    if not word in hash_dict:

                        # initialize hash object
                        h = hashlib.sha256()

                        # encrypt word type
                        h.update(word.encode('utf-8'))
                        hash_dict[word] = h.hexdigest()[:3]

                    # append encryted word
                    encrypted_tokens.append(hash_dict[word])

                speech_segment['encrypted_text'] = encrypted_tokens
                speech_segment.pop('text', None)
                new_speech_segments.append(speech_segment)

            # update annotation file
            annotations['seasons'][i]['episodes'][j]['data']['speech_segments'] = new_speech_segments

    # write out annotation file with encrypted text
    with open(output_file, 'w') as outfile:
        json.dump(annotations, outfile, indent=2)

    print('# word types: {}'.format(len(hash_dict)))
    print('# hash types: {}'.format(len(set(hash_dict.values()))))

def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--annot_file',
                        type=str,
                        help='Annotation file.',
                        required=True)

    parser.add_argument('--output_file',
                        type=str,
                        help='Output file.',
                        required=True)

    return parser.parse_args(argv)

if __name__ == '__main__':
    args = parse_arguments(sys.argv[1:])
    encrypt_text(args.annot_file, args.output_file)