lidp commited on
Commit
edccacf
·
1 Parent(s): f2ca0b5

Fix docx images (#2756)

Browse files

### What problem does this PR solve?

#2755

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

Files changed (1) hide show
  1. rag/app/naive.py +18 -9
rag/app/naive.py CHANGED
@@ -16,14 +16,15 @@ from docx import Document
16
  from timeit import default_timer as timer
17
  import re
18
  from deepdoc.parser.pdf_parser import PlainParser
19
- from rag.nlp import rag_tokenizer, naive_merge, tokenize_table, tokenize_chunks, find_codec, concat_img, naive_merge_docx, tokenize_chunks_docx
 
20
  from deepdoc.parser import PdfParser, ExcelParser, DocxParser, HtmlParser, JsonParser, MarkdownParser, TxtParser
21
  from rag.settings import cron_logger
22
  from rag.utils import num_tokens_from_string
23
  from PIL import Image
24
  from functools import reduce
25
  from markdown import markdown
26
- from docx.image.exceptions import UnrecognizedImageError
27
 
28
 
29
  class Docx(DocxParser):
@@ -42,6 +43,12 @@ class Docx(DocxParser):
42
  except UnrecognizedImageError:
43
  print("Unrecognized image format. Skipping image.")
44
  return None
 
 
 
 
 
 
45
  try:
46
  image = Image.open(BytesIO(image_blob)).convert('RGB')
47
  return image
@@ -101,7 +108,7 @@ class Docx(DocxParser):
101
  while i < len(r.cells):
102
  span = 1
103
  c = r.cells[i]
104
- for j in range(i+1, len(r.cells)):
105
  if c.text == r.cells[j].text:
106
  span += 1
107
  i = j
@@ -136,9 +143,9 @@ class Pdf(PdfParser):
136
  self._text_merge()
137
  callback(0.67, "Text merging finished")
138
  tbls = self._extract_table_figure(True, zoomin, True, True)
139
- #self._naive_vertical_merge()
140
  self._concat_downward()
141
- #self._filter_forpages()
142
 
143
  cron_logger.info("layouts: {}".format(timer() - start))
144
  return [(b["text"], self._line_tag(b, zoomin))
@@ -158,8 +165,8 @@ class Markdown(MarkdownParser):
158
  tbls = []
159
  for sec in remainder.split("\n"):
160
  if num_tokens_from_string(sec) > 10 * self.chunk_token_num:
161
- sections.append((sec[:int(len(sec)/2)], ""))
162
- sections.append((sec[int(len(sec)/2):], ""))
163
  else:
164
  sections.append((sec, ""))
165
  print(tables)
@@ -191,7 +198,7 @@ def chunk(filename, binary=None, from_page=0, to_page=100000,
191
  if re.search(r"\.docx$", filename, re.IGNORECASE):
192
  callback(0.1, "Start to parse.")
193
  sections, tbls = Docx()(filename, binary)
194
- res = tokenize_table(tbls, doc, eng) # just for table
195
 
196
  callback(0.8, "Finish parsing.")
197
  st = timer()
@@ -229,7 +236,7 @@ def chunk(filename, binary=None, from_page=0, to_page=100000,
229
  parser_config.get("chunk_token_num", 128),
230
  parser_config.get("delimiter", "\n!?;。;!?"))
231
  callback(0.8, "Finish parsing.")
232
-
233
  elif re.search(r"\.(md|markdown)$", filename, re.IGNORECASE):
234
  callback(0.1, "Start to parse.")
235
  sections, tbls = Markdown(int(parser_config.get("chunk_token_num", 128)))(filename, binary)
@@ -276,7 +283,9 @@ def chunk(filename, binary=None, from_page=0, to_page=100000,
276
  if __name__ == "__main__":
277
  import sys
278
 
 
279
  def dummy(prog=None, msg=""):
280
  pass
281
 
 
282
  chunk(sys.argv[1], from_page=0, to_page=10, callback=dummy)
 
16
  from timeit import default_timer as timer
17
  import re
18
  from deepdoc.parser.pdf_parser import PlainParser
19
+ from rag.nlp import rag_tokenizer, naive_merge, tokenize_table, tokenize_chunks, find_codec, concat_img, \
20
+ naive_merge_docx, tokenize_chunks_docx
21
  from deepdoc.parser import PdfParser, ExcelParser, DocxParser, HtmlParser, JsonParser, MarkdownParser, TxtParser
22
  from rag.settings import cron_logger
23
  from rag.utils import num_tokens_from_string
24
  from PIL import Image
25
  from functools import reduce
26
  from markdown import markdown
27
+ from docx.image.exceptions import UnrecognizedImageError, UnexpectedEndOfFileError, InvalidImageStreamError
28
 
29
 
30
  class Docx(DocxParser):
 
43
  except UnrecognizedImageError:
44
  print("Unrecognized image format. Skipping image.")
45
  return None
46
+ except UnexpectedEndOfFileError:
47
+ print("EOF was unexpectedly encountered while reading an image stream. Skipping image.")
48
+ return None
49
+ except InvalidImageStreamError:
50
+ print("The recognized image stream appears to be corrupted. Skipping image.")
51
+ return None
52
  try:
53
  image = Image.open(BytesIO(image_blob)).convert('RGB')
54
  return image
 
108
  while i < len(r.cells):
109
  span = 1
110
  c = r.cells[i]
111
+ for j in range(i + 1, len(r.cells)):
112
  if c.text == r.cells[j].text:
113
  span += 1
114
  i = j
 
143
  self._text_merge()
144
  callback(0.67, "Text merging finished")
145
  tbls = self._extract_table_figure(True, zoomin, True, True)
146
+ # self._naive_vertical_merge()
147
  self._concat_downward()
148
+ # self._filter_forpages()
149
 
150
  cron_logger.info("layouts: {}".format(timer() - start))
151
  return [(b["text"], self._line_tag(b, zoomin))
 
165
  tbls = []
166
  for sec in remainder.split("\n"):
167
  if num_tokens_from_string(sec) > 10 * self.chunk_token_num:
168
+ sections.append((sec[:int(len(sec) / 2)], ""))
169
+ sections.append((sec[int(len(sec) / 2):], ""))
170
  else:
171
  sections.append((sec, ""))
172
  print(tables)
 
198
  if re.search(r"\.docx$", filename, re.IGNORECASE):
199
  callback(0.1, "Start to parse.")
200
  sections, tbls = Docx()(filename, binary)
201
+ res = tokenize_table(tbls, doc, eng) # just for table
202
 
203
  callback(0.8, "Finish parsing.")
204
  st = timer()
 
236
  parser_config.get("chunk_token_num", 128),
237
  parser_config.get("delimiter", "\n!?;。;!?"))
238
  callback(0.8, "Finish parsing.")
239
+
240
  elif re.search(r"\.(md|markdown)$", filename, re.IGNORECASE):
241
  callback(0.1, "Start to parse.")
242
  sections, tbls = Markdown(int(parser_config.get("chunk_token_num", 128)))(filename, binary)
 
283
  if __name__ == "__main__":
284
  import sys
285
 
286
+
287
  def dummy(prog=None, msg=""):
288
  pass
289
 
290
+
291
  chunk(sys.argv[1], from_page=0, to_page=10, callback=dummy)