Added langEngines. Renamed previous method of generating alt-text as V1. Integrated langEngines to refine results of V1 Dataflow. Created V2 Dataflow. Added pydocstrings to all abstract functions.
parent
12ec75e996
commit
7541877ec7
|
@ -6,115 +6,359 @@ import bs4
|
|||
import ebooklib
|
||||
from ebooklib import epub
|
||||
|
||||
from .ocrengine import OCREngine
|
||||
from .descengine import DescEngine
|
||||
from .ocrengine import OCREngine
|
||||
from .langengine import LangEngine
|
||||
|
||||
|
||||
DEFOPTIONS = {
|
||||
"withContext": True,
|
||||
"withHash": True,
|
||||
"multiThreaded": True,
|
||||
"version": 2,
|
||||
}
|
||||
|
||||
|
||||
### ALTTEXT CLASSES
|
||||
class AltText(ABC):
|
||||
# PARSING METHODS
|
||||
@abstractmethod
|
||||
def checkData(self) -> bool:
|
||||
def setDescEngine(self, descEngine: DescEngine) -> bool:
|
||||
"""Sets current description engine.
|
||||
|
||||
Args:
|
||||
descEngine (DescEngine): A description engine.
|
||||
|
||||
Returns:
|
||||
bool: True if successful.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def parse(self, html: str) -> bs4.BeautifulSoup | epub.EpubBook:
|
||||
def setOCREngine(self, ocrEngine: OCREngine) -> bool:
|
||||
"""Sets current OCR engine.
|
||||
|
||||
Args:
|
||||
ocrEngine (OCREngine): An OCR engine.
|
||||
|
||||
Returns:
|
||||
bool: True if successful.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def setLangEngine(self, langEngine: LangEngine) -> bool:
|
||||
"""Sets current language engine.
|
||||
|
||||
Args:
|
||||
langEngine (LangEngine): A language engine.
|
||||
|
||||
Returns:
|
||||
bool: True if successful.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def setOptions(self, options: dict) -> bool:
|
||||
"""Sets current options.
|
||||
|
||||
Args:
|
||||
options (dict): A subset of DEFOPTIONS. See DEFOPTIONS constant for possible fields.
|
||||
|
||||
Returns:
|
||||
bool: True if successful.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def checkData(self) -> bool:
|
||||
"""Checks if current data exists.
|
||||
|
||||
Returns:
|
||||
bool: True if data exists.
|
||||
|
||||
Raises:
|
||||
Exception: If no data exists.
|
||||
"""
|
||||
pass
|
||||
|
||||
# PARSING METHODS
|
||||
@abstractmethod
|
||||
def parse(self, data: str) -> bs4.BeautifulSoup | epub.EpubBook:
|
||||
"""Parses data into a BeautifulSoup or EpubBook object.
|
||||
|
||||
Args:
|
||||
data (str): HTML or EPUB data.
|
||||
|
||||
Returns:
|
||||
bs4.BeautifulSoup | epub.EpubBook: The BeautifulSoup or EpubBook object stored in self.data.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def parseFile(self, filepath: str) -> bs4.BeautifulSoup | epub.EpubBook:
|
||||
"""Parses data from a file into a BeautifulSoup or EpubBook object.
|
||||
|
||||
Args:
|
||||
filepath (str): Path to HTML or EPUB file.
|
||||
|
||||
Returns:
|
||||
bs4.BeautifulSoup | epub.EpubBook: The BeautifulSoup or EpubBook object stored in self.data.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def getAllImgs(self) -> typing.List[bs4.element.Tag]:
|
||||
"""Gets all img tags.
|
||||
|
||||
Returns:
|
||||
typing.List[bs4.element.Tag]: A list of img tags.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def getNoAltImgs(self) -> typing.List[bs4.element.Tag]:
|
||||
"""Gets all img tags that either do not have an alt attribute or alt.strip() is an empty string.
|
||||
|
||||
Returns:
|
||||
typing.List[bs4.element.Tag]: A list of img tags.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def getImg(self, src: str) -> bs4.element.Tag:
|
||||
"""Gets an img tag given a src.
|
||||
|
||||
Args:
|
||||
src (str): Image source.
|
||||
|
||||
Returns:
|
||||
bs4.element.Tag: An img tag.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def setAlt(self, src: str, text: str) -> bs4.element.Tag:
|
||||
"""Sets the alt of an img tag given a src.
|
||||
|
||||
Args:
|
||||
src (str): Image source.
|
||||
text (str): New alt-text.
|
||||
|
||||
Returns:
|
||||
bs4.element.Tag: Newly modified img tag.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def setAlts(self, associations: list[dict]) -> list[bs4.element.Tag]:
|
||||
"""Sets the alt of multiple img tags given a list of associations.
|
||||
|
||||
Args:
|
||||
associations (list[dict]): A list of associations. Must have keys "src" and "alt".
|
||||
|
||||
Returns:
|
||||
list[bs4.element.Tag]: A list of newly modified img tags.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def export(self) -> str | epub.EpubBook:
|
||||
"""Exports the current data.
|
||||
|
||||
Returns:
|
||||
str | epub.EpubBook: A string of HTML or an epub.EpubBook object.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def exportToFile(self, path: str) -> str:
|
||||
"""Exports the current data to a file.
|
||||
|
||||
Args:
|
||||
path (str): A path to the file to be written.
|
||||
|
||||
Returns:
|
||||
str: The path to the file written.
|
||||
"""
|
||||
pass
|
||||
|
||||
# GENERATIVE METHODS
|
||||
@abstractmethod
|
||||
def getImgData(self, src: str) -> bytes:
|
||||
def ingest(self) -> bool:
|
||||
"""Uploads the current data and to the language engine for ingestion.
|
||||
This allows the language engine to reference the current data as a document.
|
||||
|
||||
Returns:
|
||||
bool: True if successful.
|
||||
|
||||
Raises:
|
||||
Exception: If no langEngine is set.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def getContext(self, tag: bs4.Tag) -> str:
|
||||
def degest(self) -> bool:
|
||||
"""Deletes the current data from the language engine.
|
||||
|
||||
Returns:
|
||||
bool: True if successful.
|
||||
|
||||
Raises:
|
||||
Exception: If no langEngine is set.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def getImgData(self, src: str) -> bytes:
|
||||
"""Gets byte data of an image given a src.
|
||||
|
||||
Args:
|
||||
src (str): Image source.
|
||||
|
||||
Returns:
|
||||
bytes: Image data as bytes.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def getContext(self, tag: bs4.Tag) -> list[str]:
|
||||
"""Gets the context of an img tag.
|
||||
Context being the text immediately before and after the img tag.
|
||||
|
||||
Args:
|
||||
tag (bs4.Tag): The img tag to get context for.
|
||||
|
||||
Returns:
|
||||
list[str]: A list of length 2. The first element is the text immediately before the img tag. The second element is the text immediately after the img tag.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def genChars(self, imgData: bytes, src: str) -> str:
|
||||
"""Searches for characters in an image.
|
||||
|
||||
Args:
|
||||
imgData (bytes): Image data as bytes.
|
||||
src (str): Source of the image.
|
||||
|
||||
Returns:
|
||||
str: String of characters found in the image.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def genDesc(self, imgData: bytes, src: str, context: str = None) -> str:
|
||||
"""Generates a description of an image.
|
||||
|
||||
Args:
|
||||
imgData (bytes): Image data as bytes.
|
||||
src (str): Source of the image.
|
||||
context (str, optional): Context for an image. See getContext for more information. Defaults to None.
|
||||
|
||||
Returns:
|
||||
str: Description of the image.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def genAltText(
|
||||
self, src: str, withcontext: bool = False, withocr: bool = False
|
||||
) -> str:
|
||||
def genAltTextV1(self, src: str) -> str:
|
||||
"""Generates alt-text for an image given its source.
|
||||
Uses V1 Dataflow model. This means the description and characters are generated and optionally refined separately.
|
||||
|
||||
Args:
|
||||
src (str): Source of the image.
|
||||
|
||||
Returns:
|
||||
str: Generated alt-text for the image.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def genAltTextV2(self, src: str) -> str:
|
||||
"""Generates alt-text for an image given its source.
|
||||
Uses V2 Dataflow model. This means the description and characters are generated and then alt-text is generated using both pieces of information.
|
||||
|
||||
Args:
|
||||
src (str): Source of the image.
|
||||
|
||||
Returns:
|
||||
str: Generated alt-text for the image.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def genAltText(self, src: str) -> str:
|
||||
"""Generates alt-text for an image given its source and current options.
|
||||
|
||||
Args:
|
||||
src (str): Source of the image.
|
||||
|
||||
Returns:
|
||||
str: Generated alt-text for the image.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def genAssociation(
|
||||
self,
|
||||
tag: bs4.element.Tag,
|
||||
withcontext: bool = False,
|
||||
withocr: bool = False,
|
||||
withhash: bool = False,
|
||||
) -> dict:
|
||||
"""Generates alt-text and returns an association given an img tag and current options.
|
||||
|
||||
Args:
|
||||
tag (bs4.element.Tag): Image tag to make an association for.
|
||||
|
||||
Returns:
|
||||
dict: The association. Must have keys "src" and "alt". If "withHash" is True, must also have key "hash".
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _genAltAssociationsST(
|
||||
self,
|
||||
tags: list[bs4.element.Tag],
|
||||
withcontext: bool = False,
|
||||
withocr: bool = False,
|
||||
withhash: bool = False,
|
||||
) -> list[dict]:
|
||||
"""Generates alt-text and creates associations given a list of img tags and current options.
|
||||
Single threaded implementation.
|
||||
|
||||
Args:
|
||||
tags (list[bs4.element.Tag]): List of img tags to make associations for.
|
||||
|
||||
Returns:
|
||||
list[dict]: List of associations. Must have keys "src" and "alt". If "withHash" is True, must also have key "hash".
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _genAltAssociationsMT(
|
||||
self,
|
||||
tags: list[bs4.element.Tag],
|
||||
withcontext: bool = False,
|
||||
withocr: bool = False,
|
||||
withhash: bool = False,
|
||||
) -> list[dict]:
|
||||
"""Generates alt-text and creates associations given a list of img tags and current options.
|
||||
Multi threaded implementation.
|
||||
|
||||
Args:
|
||||
tags (list[bs4.element.Tag]): List of img tags to make associations for.
|
||||
|
||||
Returns:
|
||||
list[dict]: List of associations. Must have keys "src" and "alt". If "withHash" is True, must also have key "hash".
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def genAltAssociations(
|
||||
self,
|
||||
tags: list[bs4.element.Tag],
|
||||
withcontext: bool = False,
|
||||
withocr: bool = False,
|
||||
withhash: bool = False,
|
||||
multithreaded: bool = True,
|
||||
) -> list[dict]:
|
||||
"""Generates alt-text and creates associations given a list of img tags and current options.
|
||||
Automatically selects mutli or single threaded implementation based on current options.
|
||||
|
||||
Args:
|
||||
tags (list[bs4.element.Tag]): List of img tags to make associations for.
|
||||
|
||||
Returns:
|
||||
list[dict]: List of associations. Must have keys "src" and "alt". If "withHash" is True, must also have key "hash".
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
|
@ -134,15 +378,43 @@ def getSoup(content: str) -> bs4.BeautifulSoup:
|
|||
### IMPLEMENTATIONS
|
||||
class AltTextHTML(AltText):
|
||||
def __init__(
|
||||
self, descEngine: DescEngine = None, ocrEngine: OCREngine = None
|
||||
self,
|
||||
descEngine: DescEngine,
|
||||
ocrEngine: OCREngine = None,
|
||||
langEngine: LangEngine = None,
|
||||
options: dict = {},
|
||||
) -> None:
|
||||
self.data = None
|
||||
self.filename = None
|
||||
self.filedir = None
|
||||
|
||||
self.descEngine = descEngine
|
||||
self.ocrEngine = ocrEngine
|
||||
self.langEngine = langEngine
|
||||
|
||||
self.options = DEFOPTIONS
|
||||
for key in dict.keys(options):
|
||||
self.options[key] = options[key]
|
||||
|
||||
return None
|
||||
|
||||
def setDescEngine(self, descEngine: DescEngine) -> bool:
|
||||
self.descEngine = descEngine
|
||||
return True
|
||||
|
||||
def setOCREngine(self, ocrEngine: OCREngine) -> bool:
|
||||
self.descEngine = ocrEngine
|
||||
return True
|
||||
|
||||
def setLangEngine(self, langEngine: LangEngine) -> bool:
|
||||
self.descEngine = langEngine
|
||||
return True
|
||||
|
||||
def setOptions(self, options: dict) -> bool:
|
||||
for key in dict.keys(options):
|
||||
self.options[key] = options[key]
|
||||
return True
|
||||
|
||||
def checkData(self) -> bool:
|
||||
if not hasattr(self, "data") or self.data == None:
|
||||
raise Exception("no data set. please use .parse or .parseFile")
|
||||
|
@ -156,6 +428,7 @@ class AltTextHTML(AltText):
|
|||
|
||||
def parseFile(self, filepath: str) -> bs4.BeautifulSoup:
|
||||
with open(filepath, encoding="utf8") as html:
|
||||
self.filepath = filepath
|
||||
l = filepath.split("/")
|
||||
self.filename = l.pop()
|
||||
self.filedir = "/".join(l) + "/"
|
||||
|
@ -204,6 +477,23 @@ class AltTextHTML(AltText):
|
|||
return path
|
||||
|
||||
# GENERATIVE METHODS
|
||||
def ingest(self) -> bool:
|
||||
if self.langEngine == None:
|
||||
raise Exception(
|
||||
"To use ingest, you must have an appropriate langEngine set."
|
||||
)
|
||||
with open(self.filepath, "rb") as html:
|
||||
self.langEngine.ingest(self.filename, html)
|
||||
return True
|
||||
|
||||
def degest(self) -> bool:
|
||||
if self.langEngine == None:
|
||||
raise Exception(
|
||||
"To use degest, you must have an appropriate langEngine set."
|
||||
)
|
||||
self.langEngine.degest(self.filename)
|
||||
return True
|
||||
|
||||
def __getImgFilePath(self, src: str) -> str:
|
||||
self.checkData()
|
||||
path = f"{self.filedir}{src}"
|
||||
|
@ -215,8 +505,32 @@ class AltTextHTML(AltText):
|
|||
bin = bin.read()
|
||||
return bin
|
||||
|
||||
def getContext(self, tag: bs4.Tag) -> str:
|
||||
raise Exception("IMPLEMENT ME")
|
||||
def getContext(self, tag: bs4.Tag) -> list[str]:
|
||||
context = [None, None]
|
||||
elem = tag
|
||||
text = ""
|
||||
try:
|
||||
text = elem.text.strip()
|
||||
while text == "":
|
||||
elem = elem.previous_element
|
||||
text = elem.text.strip()
|
||||
context[0] = text
|
||||
except:
|
||||
print("error 0")
|
||||
context[0] = None
|
||||
elem = tag
|
||||
text = ""
|
||||
try:
|
||||
text = elem.text.strip()
|
||||
while text == "":
|
||||
elem = elem.previous_element
|
||||
text = elem.text.strip()
|
||||
context[1] = text
|
||||
except:
|
||||
print("error 1")
|
||||
context[1] = None
|
||||
print(context)
|
||||
return context
|
||||
|
||||
def genChars(self, imgData: bytes, src: str) -> str:
|
||||
text = self.ocrEngine.genChars(imgData, src)
|
||||
|
@ -226,78 +540,77 @@ class AltTextHTML(AltText):
|
|||
alt = self.descEngine.genDesc(imgData, src, context)
|
||||
return alt
|
||||
|
||||
def genAltText(
|
||||
self, src: str, withcontext: bool = False, withocr: bool = False
|
||||
) -> str:
|
||||
def genAltTextV1(self, src: str) -> str:
|
||||
imgdata = self.getImgData(src)
|
||||
context = None
|
||||
if withcontext:
|
||||
if self.options["withContext"]:
|
||||
context = self.getContext(self.getImg(src))
|
||||
desc = self.genDesc(imgdata, src, context)
|
||||
alt = f"AUTO-GENERATED ALT-TEXT: {desc}"
|
||||
if withocr:
|
||||
chars = self.genChars(imgdata, src, context)
|
||||
alt = f"{alt}\nTEXT FOUND IN IMAGE: {chars}"
|
||||
if self.langEngine != None:
|
||||
chars = self.langEngine.refineDesc(desc)
|
||||
|
||||
alt = f"IMAGE CAPTION: {desc}"
|
||||
if self.ocrEngine != None:
|
||||
chars = self.genChars(imgdata, src)
|
||||
if self.langEngine != None:
|
||||
chars = self.langEngine.refineOCR(chars)
|
||||
alt = f"{alt}\nTEXT IN IMAGE: {chars}"
|
||||
|
||||
return alt
|
||||
|
||||
def genAltTextV2(self, src: str) -> str:
|
||||
imgdata = self.getImgData(src)
|
||||
context = [None, None]
|
||||
if self.options["withContext"]:
|
||||
context = self.getContext(self.getImg(src))
|
||||
desc = self.genDesc(imgdata, src, context)
|
||||
|
||||
chars = ""
|
||||
if self.ocrEngine != None:
|
||||
chars = self.genChars(imgdata, src).strip()
|
||||
|
||||
if self.langEngine == None:
|
||||
raise Exception("To use version 2, you must have a langEngine set.")
|
||||
|
||||
return self.langEngine.refineAlt(desc, chars, context, None)
|
||||
|
||||
def genAltText(self, src: str) -> str:
|
||||
if self.options["version"] == 1:
|
||||
return self.genAltTextV1(src)
|
||||
return self.genAltTextV2(src)
|
||||
|
||||
def genAssociation(
|
||||
self,
|
||||
tag: bs4.element.Tag,
|
||||
withcontext: bool = False,
|
||||
withocr: bool = False,
|
||||
withhash: bool = False,
|
||||
) -> dict:
|
||||
src = tag.attrs["src"]
|
||||
alt = self.genAltText(src, withcontext, withocr)
|
||||
alt = self.genAltText(src)
|
||||
association = {"src": src, "alt": alt}
|
||||
if withhash:
|
||||
if self.options["withHash"]:
|
||||
data = self.getImgData(src)
|
||||
association["hash"] = hash(data)
|
||||
return association
|
||||
|
||||
def _genAltAssociationsST(
|
||||
self,
|
||||
tags: list[bs4.element.Tag],
|
||||
withcontext: bool = False,
|
||||
withocr: bool = False,
|
||||
withhash: bool = False,
|
||||
) -> list[dict]:
|
||||
def _genAltAssociationsST(self, tags: list[bs4.element.Tag]) -> list[dict]:
|
||||
associations = []
|
||||
for tag in tags:
|
||||
associations.append(
|
||||
self.genAssociation(tag, withcontext, withocr, withhash)
|
||||
)
|
||||
associations.append(self.genAssociation(tag))
|
||||
return associations
|
||||
|
||||
def _genAltAssociationsMT(
|
||||
self,
|
||||
tags: list[bs4.element.Tag],
|
||||
withcontext: bool = False,
|
||||
withocr: bool = False,
|
||||
withhash: bool = False,
|
||||
) -> list[dict]:
|
||||
associations = []
|
||||
|
||||
def genAppend(tag, withcontext, withocr, withhash):
|
||||
associations.append(
|
||||
self.genAssociation(
|
||||
tag,
|
||||
withcontext,
|
||||
withocr,
|
||||
withhash,
|
||||
)
|
||||
)
|
||||
def genAppend(tag):
|
||||
associations.append(self.genAssociation(tag))
|
||||
|
||||
threads: list[Thread] = []
|
||||
for tag in tags:
|
||||
thread = Thread(
|
||||
target=genAppend,
|
||||
args=(
|
||||
tag,
|
||||
withcontext,
|
||||
withocr,
|
||||
withhash,
|
||||
),
|
||||
args=(tag),
|
||||
)
|
||||
thread.start()
|
||||
threads.append(thread)
|
||||
|
@ -308,14 +621,10 @@ class AltTextHTML(AltText):
|
|||
def genAltAssociations(
|
||||
self,
|
||||
tags: list[bs4.element.Tag],
|
||||
withcontext: bool = False,
|
||||
withocr: bool = False,
|
||||
withhash: bool = False,
|
||||
multithreaded: bool = True,
|
||||
) -> list[dict]:
|
||||
if multithreaded:
|
||||
return self._genAltAssociationsMT(tags, withcontext, withocr, withhash)
|
||||
return self._genAltAssociationsST(tags, withcontext, withocr, withhash)
|
||||
if self.options["multiThreaded"]:
|
||||
return self._genAltAssociationsMT(tags)
|
||||
return self._genAltAssociationsST(tags)
|
||||
|
||||
|
||||
class AltTextEPUB(AltText):
|
||||
|
|
|
@ -11,9 +11,28 @@ from vertexai.vision_models import ImageTextModel, Image
|
|||
class DescEngine(ABC):
|
||||
@abstractmethod
|
||||
def genDesc(self, imgData: bytes, src: str, context: str = None) -> str:
|
||||
"""Generates description for an image.
|
||||
|
||||
Args:
|
||||
imgData (bytes): Image data in bytes.
|
||||
src (str): Source of image.
|
||||
context (str, optional): Context of image. See getContext in alttext for more information. Defaults to None.
|
||||
|
||||
Returns:
|
||||
str: _description_
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
### TEST CLASS
|
||||
class _TDescEngine(DescEngine):
|
||||
def __init__(self):
|
||||
return None
|
||||
|
||||
def genDesc(self, imgData: bytes, src: str, context: str = None) -> str:
|
||||
return f"TEST {src}"
|
||||
|
||||
|
||||
### IMPLEMENTATIONS
|
||||
class ReplicateMiniGPT4API(DescEngine):
|
||||
def __init__(self, key: str) -> None:
|
||||
|
@ -43,6 +62,30 @@ class ReplicateMiniGPT4API(DescEngine):
|
|||
return output
|
||||
|
||||
|
||||
class ReplicateClipAPI(DescEngine):
|
||||
def __init__(self, key: str) -> None:
|
||||
self.__setKey(key)
|
||||
return None
|
||||
|
||||
def __getKey(self) -> str:
|
||||
if not hasattr(self, "data") or self.key == None:
|
||||
raise Exception("no key set. please use ._setKey(key:str)")
|
||||
return self.key
|
||||
|
||||
def __setKey(self, key: str) -> bool:
|
||||
self.key = key
|
||||
os.environ["REPLICATE_API_TOKEN"] = key
|
||||
return True
|
||||
|
||||
def genDesc(self, imgData: bytes, src: str, context: str = None) -> str:
|
||||
base64_utf8_str = base64.b64encode(imgData).decode("utf-8")
|
||||
model = "rmokady/clip_prefix_caption:9a34a6339872a03f45236f114321fb51fc7aa8269d38ae0ce5334969981e4cd8"
|
||||
ext = src.split(".")[-1]
|
||||
dataurl = f"data:image/{ext};base64,{base64_utf8_str}"
|
||||
output = replicate.run(model, input={"image": dataurl})
|
||||
return output
|
||||
|
||||
|
||||
class GoogleVertexAPI(DescEngine):
|
||||
def __init__(self, project_id: str, location: str, gac_path: str) -> None:
|
||||
self.project_id = project_id
|
||||
|
@ -74,11 +117,3 @@ class GoogleVertexAPI(DescEngine):
|
|||
language="en",
|
||||
)
|
||||
return captions[0]
|
||||
|
||||
|
||||
class __FlowTest(DescEngine):
|
||||
def __init__(self):
|
||||
return None
|
||||
|
||||
def genDesc(self, imgData: bytes, src: str, context: str = None) -> str:
|
||||
return f"TEST {src}"
|
||||
|
|
|
@ -0,0 +1,223 @@
|
|||
from abc import ABC, abstractmethod
|
||||
import requests
|
||||
|
||||
|
||||
### LANGENGINE CLASSES
|
||||
class LangEngine(ABC):
|
||||
@abstractmethod
|
||||
def _completion(self, prompt: str) -> str:
|
||||
"""Sends message to language model and returns its response.
|
||||
|
||||
Args:
|
||||
prompt (str): Prompt to send to language model.
|
||||
|
||||
Returns:
|
||||
str: Response from language model.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def refineDesc(self, description: str) -> str:
|
||||
"""Refines description of an image.
|
||||
Used in V1 Dataflow.
|
||||
|
||||
Args:
|
||||
description (str): Description of an image.
|
||||
|
||||
Returns:
|
||||
str: Refinement of description.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def refineOCR(self, chars: str) -> str:
|
||||
"""Refines characters found in an image.
|
||||
Used in V1 Dataflow.
|
||||
|
||||
Args:
|
||||
chars (str): Characters found in an image.
|
||||
|
||||
Returns:
|
||||
str: Refinement of characters.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def genPrompt(self, desc: str, chars: str, context: list[str], caption: str) -> str:
|
||||
"""Generates prompt to send to language model in V2 Dataflow.
|
||||
|
||||
Args:
|
||||
desc (str): Description of an image.
|
||||
chars (str): Characters found in an image.
|
||||
context (list[str]): Context of an image. See getContext in alttext for more information.
|
||||
caption (str): Caption of an image.
|
||||
|
||||
Returns:
|
||||
str: Prompt to send to language model.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def refineAlt(
|
||||
self,
|
||||
desc: str,
|
||||
chars: str = None,
|
||||
context: list[str] = None,
|
||||
caption: str = None,
|
||||
) -> str:
|
||||
"""Generates alt-text for an image.
|
||||
Used in V2 Dataflow.
|
||||
|
||||
Args:
|
||||
desc (str): Description of an image.
|
||||
chars (str, optional): Characters found in an image. Defaults to None.
|
||||
context (list[str], optional): Context of an image. See getContext in alttext for more information. Defaults to None.
|
||||
caption (str, optional): Caption of an image. Defaults to None.
|
||||
|
||||
Returns:
|
||||
str: Alt-text for an image.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def ingest(self, filename: str, binary) -> bool:
|
||||
"""Ingests a file into the language model.
|
||||
|
||||
Args:
|
||||
filename (str): Name of file.
|
||||
binary (_type_): Data of file.
|
||||
|
||||
Returns:
|
||||
bool: True if successful.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def degest(self, filename: str) -> bool:
|
||||
"""Removes a file from the language model.
|
||||
|
||||
Args:
|
||||
filename (str): Name of file.
|
||||
|
||||
Returns:
|
||||
bool: True if successful.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
### IMPLEMENTATIONS
|
||||
class PrivateGPT(LangEngine):
|
||||
def __init__(self, host) -> None:
|
||||
self.host = host
|
||||
|
||||
def __setHost(self, host) -> bool:
|
||||
self.host = host
|
||||
return True
|
||||
|
||||
def _completion(self, prompt: str) -> str:
|
||||
body = {
|
||||
"include_sources": False,
|
||||
"prompt": prompt,
|
||||
"stream": False,
|
||||
"use_context": True,
|
||||
}
|
||||
r = requests.post(f"{self.host}/v1/completions", json=body)
|
||||
r = r.json()
|
||||
return r["choices"][0]["message"]["content"].strip()
|
||||
|
||||
def refineDesc(self, description: str) -> str:
|
||||
prompt = f"""The following string surrounded with '///' was generated by an Image Captioning AI when ran on some arbitrary image.
|
||||
///{description}///
|
||||
|
||||
Your goal is to refine the string to be inserted as alt-text for an image in an Ebook.
|
||||
|
||||
Here are guidelines to follow...
|
||||
1. Prioritize information in text alternative:
|
||||
Aim to put the most important information at the beginning.
|
||||
2. Length of the text alternative:
|
||||
The alt text should be the most concise description possible of the image's purpose. If anything more than a short phrase or sentence is needed, it would be better to use one of the long description methods discussed in complex images.
|
||||
3. Superfluous information in the text alternative:
|
||||
Usually, there's no need to include words like “image”, “icon”, or “picture” in the alt text. People who can see will know this already, and screen readers announce the presence of an image. In some situations, it may be important to distinguish between paintings, photographs, or illustrations, etc., but it's best to avoid the more generic use of the terms.
|
||||
|
||||
Format your response as...
|
||||
The refined string is: <refined_string>
|
||||
|
||||
If the string is empty, simply respond with...
|
||||
The refined string is: N/A"""
|
||||
return self._completion(prompt)
|
||||
|
||||
def refineOCR(self, chars: str) -> str:
|
||||
prompt = f"""The following string surrounded with '///' was generated by an Optical Character Recognition software when ran on some arbitrary image.
|
||||
///
|
||||
{chars}
|
||||
///
|
||||
|
||||
Your goal is to refine the string.
|
||||
There may be random/excess spaces or other characters in the string, please remove them.
|
||||
Do not surround the refined string in quotation marks.
|
||||
|
||||
Format your response as...
|
||||
The refined string is: <refined_string>
|
||||
|
||||
If the string is empty, simply respond with...
|
||||
The refined string is: N/A"""
|
||||
return self._completion(prompt)
|
||||
|
||||
def genPrompt(self, desc: str, chars: str, context: list[str], caption: str) -> str:
|
||||
ocr = ""
|
||||
if chars != None and chars != "":
|
||||
ocr = f"\nThe following string surrounded with '///' was generated by an Optical Character Recognition software when ran on the image.\n///{chars}///"
|
||||
before = ""
|
||||
if context[0] != None and context[0] != "":
|
||||
before = f"\nThe following string surrounded with '///' is the nearest text found before the image.\n///{context[0]}///"
|
||||
after = ""
|
||||
if context[1] != None and context[1] != "":
|
||||
after = f"\nThe following string surrounded with '///' is the nearest text found after the image.\n///{context[1]}///"
|
||||
cap = ""
|
||||
if caption != None and caption != "":
|
||||
cap = f"\nThe following string surrounded with '///' is a caption in the Ebook for the image.\n///{caption}///"
|
||||
|
||||
prompt = f"""There following information is regarding an image found in an Ebook with no alternative-text.
|
||||
The following string surrounded with '///' was generated by an Image Captioning AI when ran on the image.
|
||||
///{desc}///{ocr}{cap}{before}{after}
|
||||
|
||||
Your goal is to create alternative-text for the image given the prior information.
|
||||
|
||||
Here are guidelines to follow to create quality alt-text...
|
||||
1. Prioritize information in text alternative:
|
||||
Aim to put the most important information at the beginning.
|
||||
2. Length of the text alternative:
|
||||
The alt text should be the most concise description possible of the image's purpose. If anything more than a short phrase or sentence is needed, it would be better to use one of the long description methods discussed in complex images.
|
||||
3. Superfluous information in the text alternative:
|
||||
Usually, there's no need to include words like “image”, “icon”, or “picture” in the alt text. People who can see will know this already, and screen readers announce the presence of an image. In some situations, it may be important to distinguish between paintings, photographs, or illustrations, etc., but it's best to avoid the more generic use of the terms.
|
||||
|
||||
Using all of the information stated, please generate alt-text for the image.
|
||||
In your response, please only give the alt-text."""
|
||||
return prompt
|
||||
|
||||
def refineAlt(
|
||||
self,
|
||||
desc: str,
|
||||
chars: str = None,
|
||||
context: list[str] = None,
|
||||
caption: str = None,
|
||||
) -> str:
|
||||
prompt = self.genPrompt(
|
||||
desc,
|
||||
chars,
|
||||
context,
|
||||
caption,
|
||||
)
|
||||
return self._completion(prompt)
|
||||
|
||||
def ingest(self, filename: str, binary) -> bool:
|
||||
ext = filename.split(".")[1]
|
||||
files = {"file": (filename, binary, f"application/{ext}")}
|
||||
headers = {"accept": "application/json"}
|
||||
r = requests.post(f"{self.host}/v1/ingest", files=files, headers=headers)
|
||||
return True
|
||||
|
||||
def degest(self, filename: str) -> bool:
|
||||
headers = {"accept": "application/json"}
|
||||
r = requests.delete(f"{self.host}/v1/ingest/{filename}", headers=headers)
|
||||
return True
|
|
@ -5,13 +5,29 @@ from io import BytesIO
|
|||
import pytesseract
|
||||
|
||||
|
||||
### OCRENGINE CLASSES
|
||||
### OCRENGINE ABSTRACT
|
||||
class OCREngine(ABC):
|
||||
@abstractmethod
|
||||
def genChars(self, imgData: bytes, src: str, context: str = None) -> str:
|
||||
"""Searches for characters in an image.
|
||||
|
||||
Args:
|
||||
imgData (bytes): Image data in bytes.
|
||||
src (str): Image source.
|
||||
context (str, optional): Context of an image. See getContext in alttext for more information. Defaults to None.
|
||||
|
||||
Returns:
|
||||
str: Characters found in an image.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
### TEST CLASS
|
||||
class _TOCREngine(OCREngine):
|
||||
def genChars(self, imgData: bytes, src: str, context: str = None) -> str:
|
||||
return f"TEST {src}"
|
||||
|
||||
|
||||
### IMPLEMENTATIONS
|
||||
class Tesseract(OCREngine):
|
||||
def __init__(self) -> None:
|
||||
|
|
|
@ -4,6 +4,7 @@ sys.path.append("../")
|
|||
import src.alttext.alttext as alttext
|
||||
import src.alttext.descengine as descengine
|
||||
import src.alttext.ocrengine as ocrengine
|
||||
import src.alttext.langengine as langengine
|
||||
import keys
|
||||
|
||||
# HTML BOOK FILEPATHS
|
||||
|
@ -17,33 +18,23 @@ EPUB1 = "../books/pg71856-images-3.epub"
|
|||
EPUB2 = "../books/pg71908-images-3.epub"
|
||||
EPUB3 = "../books/seuss.epub"
|
||||
|
||||
HOST1 = "http://127.0.0.1:8001"
|
||||
|
||||
|
||||
def testHTML():
|
||||
print("TESTING HTML")
|
||||
|
||||
alt: alttext.AltTextHTML = alttext.AltTextHTML(
|
||||
descengine.GoogleVertexAPI(
|
||||
keys.VertexProject(), keys.VertexRegion(), keys.VertexGAC()
|
||||
),
|
||||
# descengine.GoogleVertexAPI(
|
||||
# keys.VertexProject(), keys.VertexRegion(), keys.VertexGAC()
|
||||
# ),
|
||||
# descengine.ReplicateMiniGPT4API(keys.ReplicateEricKey()),
|
||||
descengine.ReplicateClipAPI(keys.ReplicateEricKey()),
|
||||
ocrengine.Tesseract(),
|
||||
langengine.PrivateGPT(HOST1),
|
||||
)
|
||||
alt.parseFile(HTML_INFINITY)
|
||||
|
||||
imgs = alt.getAllImgs()
|
||||
print(imgs)
|
||||
l = alt.genAltAssociations(imgs, False, False, True, True)
|
||||
print(l)
|
||||
|
||||
|
||||
def testEPUB():
|
||||
print("TESTING EPUB")
|
||||
altEPUB: alttext.AltTextEPUB = alttext.AltTextEPUB()
|
||||
|
||||
altEPUB.parseFile(EPUB2)
|
||||
imgs = altEPUB.getNoAltImgs()
|
||||
print(imgs)
|
||||
alt.parseFile(HTML_HUNTING)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
testHTML()
|
||||
# testEPUB()
|
||||
|
|
Loading…
Reference in New Issue