Fixed bug in multi-threaded association creation. Removed test engines from ocrengine and descengine. Condensed Replicate DescEngines into one.

.gitignore

@@ -2,5 +2,7 @@
 **/__pycache__/
 
 /books
+/tests/outputs
+
 **/keys.py
 **/vertex-key.json

src/alttext/alttext.py

@@ -610,7 +610,7 @@ class AltTextHTML(AltText):
         for tag in tags:
             thread = Thread(
                 target=genAppend,
-                args=(tag),
+                args=(tag,),
             )
             thread.start()
             threads.append(thread)

src/alttext/descengine.py

@@ -24,34 +24,45 @@ class DescEngine(ABC):
         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}"
+REPLICATE_MODELS = {
+    "blip": "salesforce/blip:2e1dddc8621f72155f24cf2e0adbde548458d3cab9f00c0139eea840d0ac4746",
+    "clip_prefix_caption": "rmokady/clip_prefix_caption:9a34a6339872a03f45236f114321fb51fc7aa8269d38ae0ce5334969981e4cd8",
+    "clip-caption-reward": "j-min/clip-caption-reward:de37751f75135f7ebbe62548e27d6740d5155dfefdf6447db35c9865253d7e06",
+    "img2prompt": "methexis-inc/img2prompt:50adaf2d3ad20a6f911a8a9e3ccf777b263b8596fbd2c8fc26e8888f8a0edbb5",
+    "minigpt4": "daanelson/minigpt-4:b96a2f33cc8e4b0aa23eacfce731b9c41a7d9466d9ed4e167375587b54db9423",
+    "image-captioning-with-visual-attention": "nohamoamary/image-captioning-with-visual-attention:9bb60a6baa58801aa7cd4c4fafc95fcf1531bf59b84962aff5a718f4d1f58986",
+}
 
 
 ### IMPLEMENTATIONS
-class ReplicateMiniGPT4API(DescEngine):
-    def __init__(self, key: str) -> None:
+class ReplicateAPI(DescEngine):
+    def __init__(self, key: str, model: str = "blip") -> None:
         self.__setKey(key)
+        self.__setModel(model)
         return None
 
+    def __getModel(self) -> str:
+        return self.model
+
+    def __setModel(self, modelName: str) -> str:
+        if modelName not in REPLICATE_MODELS:
+            raise Exception(
+                f"{modelName} is not a valid model. Please choose from {list(REPLICATE_MODELS.keys())}"
+            )
+        self.model = REPLICATE_MODELS[modelName]
+        return self.model
+
     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:
+    def __setKey(self, key: str) -> str:
         self.key = key
         os.environ["REPLICATE_API_TOKEN"] = key
-        return True
+        return self.key
 
     def genDesc(self, imgData: bytes, src: str, context: str = None) -> str:
         base64_utf8_str = base64.b64encode(imgData).decode("utf-8")
-        model = "daanelson/minigpt-4:b96a2f33cc8e4b0aa23eacfce731b9c41a7d9466d9ed4e167375587b54db9423"
+        model = self.__getModel()
         ext = src.split(".")[-1]
         prompt = "Create alternative-text for this image."
         if context != None:
@@ -62,37 +73,13 @@ 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
         self.location = location
-        self.gac_path = gac_path
-
         vertexai.init(project=self.project_id, location=self.location)
+
+        self.gac_path = gac_path
         os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = self.gac_path
         return None
 

src/alttext/ocrengine.py

@@ -22,12 +22,6 @@ class OCREngine(ABC):
         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:

tests/test.py

@@ -28,12 +28,20 @@ def testHTML():
         # descengine.GoogleVertexAPI(
         #     keys.VertexProject(), keys.VertexRegion(), keys.VertexGAC()
         # ),
-        # descengine.ReplicateMiniGPT4API(keys.ReplicateEricKey()),
-        descengine.ReplicateClipAPI(keys.ReplicateEricKey()),
-        ocrengine.Tesseract(),
-        langengine.PrivateGPT(HOST1),
+        descengine.ReplicateAPI(keys.ReplicateEricKey(), "blip"),
+        # ocrengine.Tesseract(),
+        # langengine.PrivateGPT(HOST1),
+        options={"version": 1},
     )
     alt.parseFile(HTML_HUNTING)
+    imgs = alt.getAllImgs()
+    # src = imgs[5].attrs["src"]
+    # print(src)
+
+    # desc = alt.genDesc(alt.getImgData(src), src)
+    # print(desc)
+    associations = alt.genAltAssociations(imgs)
+    print(associations)
 
 
 if __name__ == "__main__":