OCR class changes

refuel-ai · Nov 11, 2024 · c767232 · c767232
1 parent 02cad6d
commit c767232
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/src/autolabel/transforms/__init__.py b/src/autolabel/transforms/__init__.py
@@ -6,6 +6,7 @@
 from .serper_api import SerperApi
 from .serper_maps import SerperMaps
 from .custom_api import CustomApi
+from .ocr import OCRTransform
 from .webpage_transform import WebpageTransform
 from .image import ImageTransform
 from typing import Dict
@@ -15,6 +16,7 @@
 logger = logging.getLogger(__name__)
 
 TRANSFORM_REGISTRY = {
+    TransformType.OCR: OCRTransform,
     TransformType.PDF: PDFTransform,
     TransformType.WEBPAGE_TRANSFORM: WebpageTransform,
     TransformType.IMAGE: ImageTransform,

diff --git a/src/autolabel/transforms/schema.py b/src/autolabel/transforms/schema.py
@@ -15,6 +15,7 @@ class TransformType(str, Enum):
     WEB_SEARCH_SERPER = "web_search"
     MAPS_SEARCH = "map_search"
     CUSTOM_API = "custom_api"
+    OCR = "ocr"
 
 
 class TransformCacheEntry(BaseModel):