Repository: mindee/doctr

Issue:
[documents] Harmonize file reading
A package user has to import different functions to read a file depending on its extension, and reading mean (from path, from bytes). This needs serious refactoring.
- [x] Handle reading means (#172)
- [x] Handle extensions (#172)


Discussion:

Resolving patch:
diff --git a/doctr/documents/reader.py b/doctr/documents/reader.py
--- a/doctr/documents/reader.py
+++ b/doctr/documents/reader.py
@@ -7,13 +7,17 @@
 import cv2
 from pathlib import Path
 import fitz
-from typing import List, Tuple, Optional, Any
+from typing import List, Tuple, Optional, Any, Union, Sequence
 
-__all__ = ['read_pdf', 'read_pdf_from_stream', 'read_img']
+__all__ = ['read_pdf', 'read_img', 'DocumentFile']
+
+
+AbstractPath = Union[str, Path]
+AbstractFile = Union[AbstractPath, bytes]
 
 
 def read_img(
-    file_path: str,
+    file: AbstractFile,
     output_size: Optional[Tuple[int, int]] = None,
     rgb_output: bool = True,
 ) -> np.ndarray:
@@ -24,17 +28,23 @@ def read_img(
         >>> page = read_img("path/to/your/doc.jpg")
 
     Args:
-        file_path: the path to the image file
+        file: the path to the image file
         output_size: the expected output size of each page in format H x W
         rgb_output: whether the output ndarray channel order should be RGB instead of BGR.
     Returns:
         the page decoded as numpy ndarray of shape H x W x 3
     """
 
-    if not Path(file_path).is_file():
-        raise FileNotFoundError(f"unable to access {file_path}")
+    if isinstance(file, (str, Path)):
+        if not Path(file).is_file():
+            raise FileNotFoundError(f"unable to access {file}")
+        img = cv2.imread(str(file), cv2.IMREAD_COLOR)
+    elif isinstance(file, bytes):
+        file = np.frombuffer(file, np.uint8)
+        img = cv2.imdecode(file, cv2.IMREAD_COLOR)
+    else:
+        raise TypeError("unsupported object type for argument 'file'")
 
-    img = cv2.imread(file_path, cv2.IMREAD_COLOR)
     # Validity check
     if img is None:
         raise ValueError("unable to read file.")
@@ -47,7 +57,7 @@ def read_img(
     return img
 
 
-def read_pdf(file_path: str, **kwargs: Any) -> List[np.ndarray]:
+def read_pdf(file: AbstractFile, **kwargs: Any) -> List[np.ndarray]:
     """Read a PDF file and convert it into an image in numpy format
 
     Example::
@@ -55,30 +65,22 @@ def read_pdf(file_path: str, **kwargs: Any) -> List[np.ndarray]:
         >>> doc = read_pdf("path/to/your/doc.pdf")
 
     Args:
-        file_path: the path to the PDF file
+        file: the path to the PDF file
     Returns:
         the list of pages decoded as numpy ndarray of shape H x W x 3
     """
 
-    # Read pages with fitz and convert them to numpy ndarrays
-    return [convert_page_to_numpy(page, **kwargs) for page in fitz.open(file_path, filetype="pdf")]
-
-
-def read_pdf_from_stream(stream: bytes, **kwargs: Any) -> List[np.ndarray]:
-    """Read a PDF stream and convert it into an image in numpy format
+    fitz_args = {}
 
-    Example::
-        >>> from doctr.documents import read_pdf_from_stream
-        >>> with open("path/to/your/doc.pdf", 'rb') as f: doc = read_pdf_from_stream(f.read())
-
-    Args:
-        stream: serialized stream of the PDF content
-    Returns:
-        the list of pages decoded as numpy ndarray of shape H x W x 3
-    """
+    if isinstance(file, (str, Path)):
+        fitz_args['filename'] = file
+    elif isinstance(file, bytes):
+        fitz_args['stream'] = file
+    else:
+        raise TypeError("unsupported object type for argument 'file'")
 
     # Read pages with fitz and convert them to numpy ndarrays
-    return [convert_page_to_numpy(page, **kwargs) for page in fitz.open(stream=stream, filetype="pdf")]
+    return [convert_page_to_numpy(page, **kwargs) for page in fitz.open(**fitz_args, filetype="pdf")]
 
 
 def convert_page_to_numpy(
@@ -114,3 +116,24 @@ def convert_page_to_numpy(
         img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
     return img
+
+
+class DocumentFile:
+    """Read a document from multiple extensions
+
+    Example::
+        >>> from doctr.documents import DocumentFile
+        >>> doc = DocumentFile.from_pdf("path/to/your/doc.pdf")
+        >>> doc = DocumentFile.from_images(["path/to/your/page1.png", "path/to/your/page2.png"])
+    """
+
+    @classmethod
+    def from_pdf(cls, file: AbstractFile, **kwargs) -> List[np.ndarray]:
+        return read_pdf(file, **kwargs)
+
+    @classmethod
+    def from_images(cls, files: Union[Sequence[AbstractFile], AbstractFile], **kwargs) -> List[np.ndarray]:
+        if isinstance(files, (str, Path, bytes)):
+            files = [files]
+
+        return [read_img(file, **kwargs) for file in files]
diff --git a/scripts/analyze.py b/scripts/analyze.py
--- a/scripts/analyze.py
+++ b/scripts/analyze.py
@@ -5,7 +5,7 @@
 
 import matplotlib.pyplot as plt
 from doctr.models import ocr_predictor
-from doctr.documents import read_pdf
+from doctr.documents import DocumentFile
 from doctr.utils.visualization import visualize_page
 
 
@@ -13,7 +13,10 @@ def main(args):
 
     model = ocr_predictor(args.detection, args.recognition, pretrained=True)
 
-    doc = read_pdf(args.path)
+    if args.path.endswith(".pdf"):
+        doc = DocumentFile.from_pdf(args.path)
+    else:
+        doc = DocumentFile.from_images(args.path)
 
     out = model(doc)
 
@@ -27,7 +30,7 @@ def parse_args():
     parser = argparse.ArgumentParser(description='DocTR end-to-end analysis',
                                      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 
-    parser.add_argument('path', type=str, help='Path to the input PDF document')
+    parser.add_argument('path', type=str, help='Path to the input document (PDF or image)')
     parser.add_argument('--detection', type=str, default='db_resnet50',
                         help='Text detection model to use for analysis')
     parser.add_argument('--recognition', type=str, default='crnn_vgg16_bn',
diff --git a/scripts/evaluate.py b/scripts/evaluate.py
--- a/scripts/evaluate.py
+++ b/scripts/evaluate.py
@@ -11,7 +11,6 @@
 
 from doctr.utils.metrics import LocalizationConfusion, ExactMatch, OCRMetric
 from doctr.datasets import FUNSD
-from doctr.documents import read_img
 from doctr.models import ocr_predictor, extract_crops
 
 


Test patch:
diff --git a/test/test_documents_reader.py b/test/test_documents_reader.py
--- a/test/test_documents_reader.py
+++ b/test/test_documents_reader.py
@@ -23,24 +23,36 @@ def test_convert_page_to_numpy(mock_pdf):
     assert resized_page.shape == (396, 306, 3)
 
 
-def _check_pdf_content(doc_tensors):
+def _check_doc_content(doc_tensors, num_pages):
     # 1 doc of 8 pages
-    assert(len(doc_tensors) == 8)
+    assert(len(doc_tensors) == num_pages)
     assert all(isinstance(page, np.ndarray) for page in doc_tensors)
     assert all(page.dtype == np.uint8 for page in doc_tensors)
 
 
-def test_read_pdf(mock_pdf):
-    doc_tensors = reader.read_pdf(mock_pdf)
-    _check_pdf_content(doc_tensors)
+def test_read_pdf(mock_pdf, mock_pdf_stream):
+    for file in [mock_pdf, mock_pdf_stream]:
+        doc_tensors = reader.read_pdf(file)
+        _check_doc_content(doc_tensors, 8)
 
+    # Wrong input type
+    with pytest.raises(TypeError):
+        _ = reader.read_pdf(123)
 
-def test_read_pdf_from_stream(mock_pdf_stream):
-    doc_tensors = reader.read_pdf_from_stream(mock_pdf_stream)
-    _check_pdf_content(doc_tensors)
 
+def test_read_img(tmpdir_factory, mock_image_stream, mock_pdf):
 
-def test_read_img(tmpdir_factory, mock_pdf):
+    # Wrong input type
+    with pytest.raises(TypeError):
+        _ = reader.read_img(123)
+
+    # Non-existing file
+    with pytest.raises(FileNotFoundError):
+        reader.read_img("my_imaginary_file.jpg")
+
+    # Invalid image
+    with pytest.raises(ValueError):
+        reader.read_img(str(mock_pdf))
 
     url = 'https://upload.wikimedia.org/wikipedia/commons/5/55/Grace_Hopper.jpg'
     file = BytesIO(requests.get(url).content)
@@ -65,9 +77,10 @@ def test_read_img(tmpdir_factory, mock_pdf):
     resized_page = reader.read_img(tmp_path, target_size)
     assert resized_page.shape[:2] == target_size
 
-    # Non-existing file
-    with pytest.raises(FileNotFoundError):
-        reader.read_img("my_imaginary_file.jpg")
-    # Invalid image
-    with pytest.raises(ValueError):
-        reader.read_img(str(mock_pdf))
+
+def test_document_file(mock_pdf, mock_image_stream, mock_image_folder):
+    pages = reader.DocumentFile.from_pdf(mock_pdf)
+    _check_doc_content(pages, 8)
+
+    pages = reader.DocumentFile.from_images(mock_image_stream)
+    _check_doc_content(pages, 1)