-
Notifications
You must be signed in to change notification settings - Fork 282
/
Ocrs.py
executable file
·336 lines (286 loc) · 11.3 KB
/
Ocrs.py
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
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
#!/usr/bin/env python3
# -*-coding:UTF-8 -*
import os
import sys
from datetime import datetime
from io import BytesIO
from PIL import Image
from PIL import ImageDraw
from pymisp import MISPObject
sys.path.append(os.environ['AIL_BIN'])
##################################
# Import Project packages
##################################
from lib.objects.abstract_daterange_object import AbstractDaterangeObject, AbstractDaterangeObjects
from lib.ConfigLoader import ConfigLoader
from packages import Date
# from lib import Language
# from lib.data_retention_engine import update_obj_date, get_obj_date_first
from flask import url_for
config_loader = ConfigLoader()
r_cache = config_loader.get_redis_conn("Redis_Cache")
r_object = config_loader.get_db_conn("Kvrocks_Objects")
baseurl = config_loader.get_config_str("Notifications", "ail_domain")
IMAGE_FOLDER = config_loader.get_files_directory('images')
config_loader = None
# SET x1,y1:x2,y2:x3,y3:x4,y4:extracted_text
class Ocr(AbstractDaterangeObject):
"""
AIL Message Object. (strings)
"""
def __init__(self, id):
super(Ocr, self).__init__('ocr', id)
def exists(self):
return r_object.exists(f'ocr:{self.id}')
def get_content(self, r_type='str'):
"""
Returns content
"""
global_id = self.get_global_id()
content = r_cache.get(f'content:{global_id}')
if not content:
dict_content = {}
for extracted in r_object.smembers(f'ocr:{self.id}'):
extracted = extracted.split(':', 4)
x, y = extracted[0].split(',', 1)
# get text line, y +- 20
rounded_y = round(int(y) / 20) * 20
if rounded_y not in dict_content:
dict_content[rounded_y] = []
dict_content[rounded_y].append((int(x), int(y), extracted[-1]))
content = ''
new_line = True
l_key = sorted(dict_content.keys())
for key in l_key:
dict_content[key] = sorted(dict_content[key], key=lambda c: c[0])
for text in dict_content[key]:
if new_line:
content = f'{content}{text[2]}'
new_line = False
else:
content = f'{content} {text[2]}'
content = f'{content}\n'
new_line = True
# Set Cache
if content:
global_id = self.get_global_id()
r_cache.set(f'content:{global_id}', content)
r_cache.expire(f'content:{global_id}', 300)
if r_type == 'str':
return content
elif r_type == 'bytes':
if content:
return content.encode()
def get_date(self): # TODO
return Date.get_today_date_str()
def get_source(self): # TODO
"""
Returns source/feeder name
"""
return 'ocr'
# l_source = self.id.split('/')[:-2]
# return os.path.join(*l_source)
def get_basename(self): # TODO
return 'ocr'
def get_language(self):
languages = self.get_languages()
if languages:
return languages.pop()
else:
return None
def get_link(self, flask_context=False):
if flask_context:
url = url_for('correlation.show_correlation', type=self.type, id=self.id)
else:
url = f'{baseurl}/correlation/show?type={self.type}&id={self.id}'
return url
def get_svg_icon(self):
return {'style': 'fas', 'icon': '\uf065', 'color': 'yellow', 'radius': 5}
def get_image_path(self):
rel_path = os.path.join(self.id[0:2], self.id[2:4], self.id[4:6], self.id[6:8], self.id[8:10], self.id[10:12], self.id[12:])
filename = os.path.join(IMAGE_FOLDER, rel_path)
return os.path.realpath(filename)
def get_misp_object(self): # TODO
obj = MISPObject('instant-message', standalone=True)
obj_date = self.get_date()
if obj_date:
obj.first_seen = obj_date
else:
self.logger.warning(
f'Export error, None seen {self.type}:{self.subtype}:{self.id}, first={obj_date}')
# obj_attrs = [obj.add_attribute('first-seen', value=obj_date),
# obj.add_attribute('raw-data', value=self.id, data=self.get_raw_content()),
# obj.add_attribute('sensor', value=get_ail_uuid())]
obj_attrs = []
for obj_attr in obj_attrs:
for tag in self.get_tags():
obj_attr.add_tag(tag)
return obj
# options: set of optional meta fields
def get_meta(self, options=None, translation_target=''):
"""
:type options: set
"""
if options is None:
options = set()
meta = self._get_meta(options=options)
meta['tags'] = self.get_tags()
meta['content'] = self.get_content()
# optional meta fields
if 'investigations' in options:
meta['investigations'] = self.get_investigations()
if 'link' in options:
meta['link'] = self.get_link(flask_context=True)
if 'icon' in options:
meta['svg_icon'] = self.get_svg_icon()
if 'img' in options:
meta['img'] = self.draw_bounding_boxs()
if 'map' in options:
meta['map'] = self.get_img_map_coords()
if 'language' in options:
meta['language'] = self.get_language()
if 'translation' in options and translation_target:
if meta.get('language'):
source = meta['language']
else:
source = None
meta['translation'] = self.translate(content=meta.get('content'), source=source, target=translation_target)
if 'language' in options:
meta['language'] = self.get_language()
return meta
def get_objs_container(self):
objs_containers = set()
# chat
objs_containers.add(self.get_first_correlation('chat'))
subchannel = self.get_first_correlation('chat-subchannel')
if subchannel:
objs_containers.add(subchannel)
thread = self.get_first_correlation('chat-thread')
if thread:
objs_containers.add(thread)
return objs_containers
def create_coord_str(self, bbox):
c1, c2, c3, c4 = bbox
x1, y1 = c1
x2, y2 = c2
x3, y3 = c3
x4, y4 = c4
return f'{int(x1)},{int(y1)}:{int(x2)},{int(y2)}:{int(x3)},{int(y3)}:{int(x4)},{int(y4)}'
def _unpack_coord(self, coord):
return coord.split(',', 1)
def get_coords(self):
coords = []
for extracted in r_object.smembers(f'ocr:{self.id}'):
coord = []
bbox = extracted.split(':', 4)[:-1]
for c in bbox:
x, y = self._unpack_coord(c)
coord.append((int(x), int(y)))
coords.append(coord)
return coords
def get_img_map_coords(self):
coords = []
for extracted in r_object.smembers(f'ocr:{self.id}'):
extract = extracted.split(':', 4)
x1, y1 = self._unpack_coord(extract[0])
x2, y2 = self._unpack_coord(extract[1])
x3, y3 = self._unpack_coord(extract[2])
x4, y4 = self._unpack_coord(extract[3])
coords.append((f'{x1},{y1},{x2},{y2},{x3},{y3},{x4},{y4}', extract[4]))
return coords
def edit_text(self, coordinates, text, new_text, new_coordinates=None):
pass
def add_text(self, coordinates, text):
val = f'{coordinates}:{text}'
return r_object.sadd(f'ocr:{self.id}', val)
def remove_text(self, val):
return r_object.srem(f'ocr:{self.id}', val)
def update_correlation(self, date=None):
if date:
self.add(date, None)
image_correl = self.get_obj_correlations('image', '', self.id)
for obj_type in image_correl:
if obj_type != 'ocr':
for obj_raw in image_correl[obj_type]:
obj_subtype, obj_id = obj_raw.split(':', 1)
self.add_correlation(obj_type, obj_subtype, obj_id)
def create(self, extracted_texts, tags=[]):
# r_object.sadd(f'{self.type}:all', self.id)
created = False
for extracted in extracted_texts:
bbox, text = extracted
if len(text) > 1:
str_coords = self.create_coord_str(bbox)
self.add_text(str_coords, text)
created = True
if created:
# Correlations
self._copy_from('image', self.id)
self.update_correlation()
self.add_correlation('image', '', self.id)
for tag in tags:
self.add_tag(tag)
return self.id
# # WARNING: UNCLEAN DELETE /!\ TEST ONLY /!\
def delete(self):
r_object.delete(f'ocr:{self.id}')
def draw_bounding_boxs(self):
img = Image.open(self.get_image_path()).convert("RGBA")
draw = ImageDraw.Draw(img)
for bbox in self.get_coords():
c1, c2, c3, c4 = bbox
draw.line((tuple(c1), tuple(c2)), fill="yellow", width=2)
draw.line((tuple(c2), tuple(c3)), fill="yellow", width=2)
draw.line((tuple(c3), tuple(c4)), fill="yellow", width=2)
draw.line((tuple(c4), tuple(c1)), fill="yellow", width=2)
# img.show()
buff = BytesIO()
img.save(buff, "PNG")
return buff.getvalue()
def create(obj_id, detections, tags=[]):
obj = Ocr(obj_id)
if not obj.exists():
obj_id = obj.create(detections, tags=tags)
if obj_id:
return obj
# TODO preload languages
def extract_text(image_path, languages, threshold=0.2):
import easyocr
reader = easyocr.Reader(languages, verbose=False)
texts = reader.readtext(image_path)
# print(texts)
extracted = []
for bbox, text, score in texts:
if score > threshold:
extracted.append((bbox, text))
return extracted
def get_ocr_languages():
return {'af', 'ar', 'as', 'az', 'be', 'bg', 'bh', 'bs', 'cs', 'cy', 'da', 'de', 'en', 'es', 'et', 'fa', 'fr', 'ga', 'hi', 'hr', 'hu', 'id', 'is', 'it', 'ja', 'kn', 'ko', 'ku', 'la', 'lt', 'lv', 'mi', 'mn', 'mr', 'ms', 'mt', 'ne', 'nl', 'no', 'oc', 'pi', 'pl', 'pt', 'ro', 'ru', 'sk', 'sl', 'sq', 'sr', 'sv', 'sw', 'ta', 'te', 'th', 'tl', 'tr', 'ug', 'uk', 'ur', 'uz', 'vi', 'zh'}
def sanityze_ocr_languages(languages, ocr_languages=None):
langs = set()
if not ocr_languages:
ocr_languages = get_ocr_languages()
for lang in languages:
if lang in ocr_languages:
if lang == 'zh':
langs.add('ch_sim')
elif lang == 'sr':
langs.add('rs_latin')
else:
langs.add(lang)
return langs
class Ocrs(AbstractDaterangeObjects):
"""
OCR Objects
"""
def __init__(self):
super().__init__('ocr', Ocr)
def sanitize_id_to_search(self, name_to_search):
return name_to_search # TODO
#### API ####
def api_get_ocr(obj_id, translation_target=None):
ocr = Ocr(obj_id)
if not ocr.exists():
return {"status": "error", "reason": "Unknown ocr"}, 404
meta = ocr.get_meta({'content', 'icon', 'img', 'language', 'link', 'map', 'translation'}, translation_target=translation_target)
return meta, 200