py311 passed

Author: Sid Mohan

datafog/processing/image_processing/donut_processor.py

@@ -7,13 +7,13 @@
 
 import requests
 from PIL import Image
+import numpy as np
 
 from .image_downloader import ImageDownloader
 
 
 class DonutProcessor:
     def __init__(self, model_path="naver-clova-ix/donut-base-finetuned-cord-v2"):
-
         self.ensure_installed("torch")
         self.ensure_installed("transformers")
 
@@ -36,13 +36,31 @@ def ensure_installed(self, package_name):
                 [sys.executable, "-m", "pip", "install", package_name]
             )
 
-    async def parse_image(self, image: Image) -> str:
+    def preprocess_image(self, image: Image.Image) -> np.ndarray:
+        # Convert to RGB if the image is not already in RGB mode
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        # Convert to numpy array
+        image_np = np.array(image)
+        
+        # Ensure the image is 3D (height, width, channels)
+        if image_np.ndim == 2:
+            image_np = np.expand_dims(image_np, axis=-1)
+            image_np = np.repeat(image_np, 3, axis=-1)
+        
+        return image_np
+
+    async def parse_image(self, image: Image.Image) -> str:
         """Process w/ DonutProcessor and VisionEncoderDecoderModel"""
+        # Preprocess the image
+        image_np = self.preprocess_image(image)
+        
         task_prompt = "<s_cord-v2>"
         decoder_input_ids = self.processor.tokenizer(
             task_prompt, add_special_tokens=False, return_tensors="pt"
         ).input_ids
-        pixel_values = self.processor(image, return_tensors="pt").pixel_values
+        pixel_values = self.processor(images=image_np, return_tensors="pt").pixel_values
 
         outputs = self.model.generate(
             pixel_values.to(self.device),
@@ -71,8 +89,8 @@ def process_url(self, url: str) -> str:
         image = self.downloader.download_image(url)
         return self.parse_image(image)
 
-    def download_image(self, url: str) -> Image:
+    def download_image(self, url: str) -> Image.Image:
         """Download an image from URL."""
         response = requests.get(url)
         image = Image.open(BytesIO(response.content))
-        return image
+        return image

datafog/services/image_service.py

@@ -4,10 +4,28 @@
 from PIL import Image
 
 from datafog.processing.image_processing.donut_processor import DonutProcessor
-from datafog.processing.image_processing.image_downloader import ImageDownloader
+# from datafog.processing.image_processing.image_downloader import ImageDownloader
 from datafog.processing.image_processing.pytesseract_processor import (
     PytesseractProcessor,
 )
+import aiohttp
+from PIL import Image
+import io
+import ssl
+import certifi
+
+class ImageDownloader:
+    async def download_image(self, url: str) -> Image.Image:
+        ssl_context = ssl.create_default_context(cafile=certifi.where())
+        async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=ssl_context)) as session:
+            async with session.get(url) as response:
+                if response.status == 200:
+                    image_data = await response.read()
+                    return Image.open(io.BytesIO(image_data))
+                else:
+                    raise Exception(f"Failed to download image. Status code: {response.status}")
+
+
 
 
 class ImageService:
@@ -20,8 +38,19 @@ def __init__(self, use_donut: bool = False, use_tesseract: bool = True):
             PytesseractProcessor() if self.use_tesseract else None
         )
 
+    # async def download_images(self, urls: List[str]) -> List[Image.Image]:
+    #     async def download_image(url: str) -> Image.Image:
+    #         return await self.downloader.download_image(url)
+
+    #     tasks = [asyncio.create_task(download_image(url)) for url in urls]
+    #     return await asyncio.gather(*tasks)
+    
     async def download_images(self, urls: List[str]) -> List[Image.Image]:
-        return await self.downloader.download_images(urls)
+        async def download_image(url: str) -> Image.Image:
+            return await self.downloader.download_image(url)
+
+        tasks = [asyncio.create_task(download_image(url)) for url in urls]
+        return await asyncio.gather(*tasks, return_exceptions=True)
 
     async def ocr_extract(
         self,

requirements-dev.txt

@@ -6,10 +6,12 @@ just
 isort
 black
 blacken-docs
+certifi
 flake8
 prettier
 tox
-pytest
+pytest==7.4.0
+pytest-asyncio==0.21.0
 pytest-cov
 mypy
 autoflake

tests/test_image_service.py

@@ -11,6 +11,7 @@
 
 
 import pytest
+import asyncio
 from PIL import Image
 
 from datafog.services.image_service import ImageService
@@ -21,13 +22,21 @@
 ]
 
 
+# @pytest.mark.asyncio
+# async def test_download_images():
+#     image_service1 = ImageService()
+#     images = await image_service1.download_images(urls)
+#     assert len(images) == 2
+#     assert all(isinstance(image, Image.Image) for image in images)
 @pytest.mark.asyncio
 async def test_download_images():
-    image_service1 = ImageService()
-    images = await image_service1.download_images(urls)
-    assert len(images) == 2
-    assert all(isinstance(image, Image.Image) for image in images)
-
+    image_service = ImageService()
+    try:
+        images = await image_service.download_images(urls)
+        assert len(images) == 2
+        assert all(isinstance(image, Image.Image) for image in images)
+    finally:
+        await asyncio.sleep(0)  # Allow pending callbacks to run
 
 @pytest.mark.asyncio
 async def test_ocr_extract_with_tesseract():

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = py312
+envlist = py311
 isolated_build = True
 
 [testenv]