File size: 5,266 Bytes
072f9dd |
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 |
import copy
import re
from io import BytesIO
from pptx import Presentation
from rag.app import callback__
from rag.nlp import huqie
from rag.parser.pdf_parser import HuParser
class Ppt(object):
def __init__(self):
super().__init__()
def __extract(self, shape):
if shape.shape_type == 19:
tb = shape.table
rows = []
for i in range(1, len(tb.rows)):
rows.append("; ".join([tb.cell(0, j).text + ": " + tb.cell(i, j).text for j in range(len(tb.columns)) if tb.cell(i, j)]))
return "\n".join(rows)
if shape.has_text_frame:
return shape.text_frame.text
if shape.shape_type == 6:
texts = []
for p in shape.shapes:
t = self.__extract(p)
if t: texts.append(t)
return "\n".join(texts)
def __call__(self, fnm, from_page, to_page, callback=None):
ppt = Presentation(fnm) if isinstance(
fnm, str) else Presentation(
BytesIO(fnm))
txts = []
self.total_page = len(ppt.slides)
for i, slide in enumerate(ppt.slides[from_page: to_page]):
texts = []
for shape in slide.shapes:
txt = self.__extract(shape)
if txt: texts.append(txt)
txts.append("\n".join(texts))
callback__((i+1)/self.total_page/2, "", callback)
callback__((min(to_page, self.total_page) - from_page) / self.total_page,
"Page {}~{}: Text extraction finished".format(from_page, min(to_page, self.total_page)), callback)
import aspose.slides as slides
import aspose.pydrawing as drawing
imgs = []
with slides.Presentation(BytesIO(fnm)) as presentation:
for i, slide in enumerate(presentation.slides[from_page: to_page]):
buffered = BytesIO()
slide.get_thumbnail(0.5, 0.5).save(buffered, drawing.imaging.ImageFormat.jpeg)
imgs.append(buffered.getvalue())
assert len(imgs) == len(txts), "Slides text and image do not match: {} vs. {}".format(len(imgs), len(txts))
callback__((min(to_page, self.total_page) - from_page) / self.total_page,
"Page {}~{}: Image extraction finished".format(from_page, min(to_page, self.total_page)), callback)
return [(txts[i], imgs[i]) for i in range(len(txts))]
class Pdf(HuParser):
def __init__(self):
super().__init__()
def __garbage(self, txt):
txt = txt.lower().strip()
if re.match(r"[0-9\.,%/-]+$", txt): return True
if len(txt) < 3:return True
return False
def __call__(self, filename, binary=None, from_page=0, to_page=100000, zoomin=3, callback=None):
self.__images__(filename if not binary else binary, zoomin, from_page, to_page)
callback__((min(to_page, self.total_page)-from_page) / self.total_page, "Page {}~{}: OCR finished".format(from_page, min(to_page, self.total_page)), callback)
assert len(self.boxes) == len(self.page_images), "{} vs. {}".format(len(self.boxes), len(self.page_images))
res = []
#################### More precisely ###################
# self._layouts_paddle(zoomin)
# self._text_merge()
# pages = {}
# for b in self.boxes:
# if self.__garbage(b["text"]):continue
# if b["page_number"] not in pages: pages[b["page_number"]] = []
# pages[b["page_number"]].append(b["text"])
# for i, lines in pages.items():
# res.append(("\n".join(lines), self.page_images[i-1]))
# return res
########################################
for i in range(len(self.boxes)):
lines = "\n".join([b["text"] for b in self.boxes[i] if not self.__garbage(b["text"])])
res.append((lines, self.page_images[i]))
return res
def chunk(filename, binary=None, from_page=0, to_page=100000, callback=None):
doc = {
"docnm_kwd": filename,
"title_tks": huqie.qie(re.sub(r"\.[a-zA-Z]+$", "", filename))
}
doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"])
res = []
if re.search(r"\.pptx?$", filename, re.IGNORECASE):
for txt,img in Ppt()(filename if not binary else binary, from_page, to_page, callback):
d = copy.deepcopy(doc)
d["content_ltks"] = huqie.qie(txt)
d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
d["image"] = img
res.append(d)
return res
if re.search(r"\.pdf$", filename, re.IGNORECASE):
for txt,img in Pdf()(filename if not binary else binary, from_page=from_page, to_page=to_page, callback=callback):
d = copy.deepcopy(doc)
d["content_ltks"] = huqie.qie(txt)
d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
d["image"] = img
res.append(d)
return res
callback__(-1, "This kind of presentation document did not support yet!", callback)
if __name__== "__main__":
import sys
print(chunk(sys.argv[1]))
|