mistral-ai-ocr 1.0__py3-none-any.whl

mistral_ai_ocr/__init__.py

@@ -0,0 +1,255 @@
+#!/usr/bin/env python
+import base64
+from typing import List
+from mistralai import Mistral, OCRImageObject
+from mistralai.models import OCRResponse
+from pathlib import Path
+import mimetypes
+from enum import Enum
+import sys
+
+class Modes(Enum):
+    FULL = 0
+    FULL_ALT = 1
+    FULL_NO_DIR = 2
+    FULL_NO_PAGES = 3
+    TEXT = 4
+    TEXT_NO_PAGES = 5
+
+def get_mode_from_string(mode_str: str):
+    for mode in Modes:
+        if mode.name == mode_str.upper() or mode.value == mode_str:
+            return mode
+    raise ValueError(f"Unknown mode: {mode_str}")
+
+def b64encode_document(document_path: Path):
+    try:
+        with open(document_path, "rb") as doc_file:
+            return base64.b64encode(doc_file.read()).decode('utf-8')
+    except FileNotFoundError:
+        return None
+    except Exception as e:
+        return None
+    
+def b64decode_document(base64_data: str, output_path: Path):
+    if ',' in base64_data:
+        _, base64_str = base64_data.split(',', 1)
+    else:
+        base64_str = base64_data
+    try:
+        image_data = base64.b64decode(base64_str)
+    except (base64.binascii.Error, ValueError) as e:
+        print(f"Error decoding base64 data: {e}", file=sys.stderr)
+        return
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(output_path, 'wb') as f:
+        f.write(image_data)
+
+class Page:
+    def __init__(self, index, markdown=None, images:List[OCRImageObject]=None):
+        self.index = index
+        self.markdown = markdown
+        self.images = images if images is not None else []
+
+    def write_markdown(self, output_path: Path, append: bool = False, insert = None):
+        if self.markdown:
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            mode = 'a' if append else 'w'
+            with open(output_path, mode) as md_file:
+                if insert:
+                    md_file.write(insert)
+                md_file.write(self.markdown)
+    
+    def write_images(self, output_directory: Path):
+        if not self.images:
+            return
+        
+        for image in self.images:
+            if image and image.image_base64:
+                image_name = image.id
+                image_path = output_directory / image_name
+                b64decode_document(image.image_base64, image_path)
+
+class MistralOCRDocument:
+    def __init__(self, 
+                 document_path: Path, 
+                 api_key: str, 
+                 include_images=True,
+                 output_directory: Path = None,
+                 generate_pages=True,
+                 full_directory_name="full",
+                 page_separator="\n",
+                 page_directory_name="page_<index>",
+                 page_text_name="<stem>.md",
+                 json_ocr_response_path=None,
+                 save_json=True,
+                ):
+        self.document_path = document_path
+        self.api_key = api_key
+        self.include_images = include_images
+        self.generate_pages = generate_pages
+        self.save_json = save_json
+        self.full_directory_name = full_directory_name
+        self.page_separator = page_separator
+        self.page_directory_name = page_directory_name
+        self.page_text_name = page_text_name
+        self.json_ocr_response_path = json_ocr_response_path
+        if output_directory is None:
+            self.output_directory = self.get_input_path().parent / self.get_input_path().stem
+        else:
+            self.output_directory = output_directory
+
+    def get_ocr_response(self, mimetype, base64_document):
+        client = Mistral(api_key=self.api_key)
+        if mimetype.startswith("image/"):
+            document_type = "image_url"
+        elif mimetype.startswith("application/pdf"):
+            document_type = "document_url"
+        else:
+            raise ValueError(f"Unsupported MIME type: {mimetype}. Only image and PDF files are supported.")
+        self.ocr_response = client.ocr.process(
+            model="mistral-ocr-latest",
+            document={
+                "type": document_type,
+                document_type: f"data:{mimetype};base64,{base64_document}" 
+            },
+            include_image_base64=self.include_images
+        )
+
+    def process_document(self):
+        if not self.document_path.exists():
+            raise FileNotFoundError(f"The document {self.document_path} does not exist.")
+        if not self.document_path.is_file():
+            raise ValueError(f"The path {self.document_path} is not a valid file.")
+        
+        mimetype, _ = mimetypes.guess_type(self.document_path)
+        if mimetype is None:
+            raise ValueError(f"Could not determine the MIME type for {self.document_path}.")
+        
+        self.get_ocr_response(mimetype, b64encode_document(self.document_path))
+        self.write_json()
+        self.process_ocr_response()
+
+    def process_json_response(self):
+        if self.json_ocr_response_path is None or not self.json_ocr_response_path.exists():
+            raise FileNotFoundError(f"The JSON OCR response {self.json_ocr_response_path} does not exist.")
+        
+        with open(self.json_ocr_response_path, "r") as json_file:
+            self.ocr_response = OCRResponse.model_validate_json(json_file.read())
+        self.write_json()
+        self.process_ocr_response()
+
+    def process(self):
+        if self.json_ocr_response_path is not None:
+            self.process_json_response()
+        else:
+            self.process_document()
+
+    def get_input_path(self):
+        if self.json_ocr_response_path is not None:
+            return self.json_ocr_response_path
+        return self.document_path
+    
+    def write_json(self):
+        if self.save_json:
+            output_path = (self.output_directory / self.get_input_path().stem).with_suffix(".json")
+            self.output_directory.mkdir(parents=True, exist_ok=True)
+            with open(output_path, "w") as text_file:
+                text_file.write(self.ocr_response.model_dump_json(indent=2))
+
+    def process_ocr_response(self):
+        response_pages = self.ocr_response.pages
+        if not response_pages:
+            print("No pages found in the OCR response.")
+            return
+        
+        pages = []
+
+        full_dir = self.output_directory / self.full_directory_name
+
+        for r_page in response_pages:
+            page = Page(
+                index=r_page.index,
+                markdown=r_page.markdown,
+                images=r_page.images
+            )
+            if self.generate_pages:
+                page_dir = self.output_directory / self.page_directory_name.replace("<index>", str(page.index))
+                page.write_markdown((
+                    page_dir / self.page_text_name.
+                    replace("<stem>", self.get_input_path().stem).
+                    replace("<index>", str(page.index))
+                    ).with_suffix(".md"))
+                if self.include_images:
+                    page.write_images(page_dir)
+            if self.include_images:
+                page.write_images(full_dir)
+            pages.append(page)
+        for i, page in enumerate(sorted(pages, key=lambda p: p.index)):
+            first = i == 0
+            md_file = (full_dir / self.get_input_path().stem).with_suffix(".md")
+            insert = self.page_separator if not first else None
+            page.write_markdown(md_file, append=not first, insert=insert)
+
+def construct_from_mode(
+    document_path: Path,
+    api_key: str,
+    output_directory: Path = None,
+    json_ocr_response_path: Path = None,
+    page_separator: str = "\n",
+    write_json: bool = True,
+    mode: Modes = Modes.FULL
+):
+    kwargs = dict(
+        document_path=document_path,
+        api_key=api_key,
+        output_directory=output_directory,
+        json_ocr_response_path=json_ocr_response_path,
+        page_separator=page_separator,
+        save_json=write_json
+    )
+    match mode:
+        case Modes.FULL:
+            kwargs.update(
+                include_images=True,
+                generate_pages=True
+            )
+        case Modes.FULL_ALT:
+            kwargs.update(
+                include_images=True,
+                generate_pages=True,
+                full_directory_name="."
+            )
+        case Modes.FULL_NO_DIR:
+            kwargs.update(
+                include_images=True,
+                generate_pages=True,
+                full_directory_name=".",
+                page_directory_name=".",
+                page_text_name="<stem><index>.md"
+            )
+        case Modes.FULL_NO_PAGES:
+            kwargs.update(
+                include_images=True,
+                generate_pages=False,
+                full_directory_name="."
+            )
+        case Modes.TEXT:
+            kwargs.update(
+                include_images=False,
+                generate_pages=True,
+                full_directory_name=".",
+                page_directory_name=".",
+                page_text_name="<stem><index>.md"
+            )
+        case Modes.TEXT_NO_PAGES:
+            kwargs.update(
+                include_images=False,
+                generate_pages=False,
+                full_directory_name="."
+            )
+        case _:
+            raise ValueError(f"Unknown mode: {mode}")
+    return MistralOCRDocument(**kwargs)

mistral_ai_ocr/__main__.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python
+from pathlib import Path
+from . import Modes, construct_from_mode, get_mode_from_string
+import argparse
+import codecs
+from os import getenv
+from dotenv import load_dotenv
+
+try:
+  load_dotenv(dotenv_path=Path.home() / ".mistral_ai_ocr.env")
+except:
+  pass
+
+_mode_choices = [mode.name for mode in Modes] + [str(mode.value) for mode in Modes]
+
+def _unescape(s: str) -> str:
+  return codecs.decode(s, 'unicode_escape')
+
+def main():
+  example_text = (
+    'examples:\n\n'
+    '%(prog)s paper.pdf\n'
+    '%(prog)s paper.pdf --api-key jrWjJE5lFketfB2sA6vvhQK2SoHQ6R39\n'
+    '%(prog)s paper.pdf -o revision\n'
+    '%(prog)s paper.pdf -e\n'
+    '%(prog)s paper.pdf -m FULL\n'
+    '%(prog)s -j paper.json\n'
+    '%(prog)s -j paper.json -m TEXT_NO_PAGES -n\n'
+  )
+  parser = argparse.ArgumentParser(
+    description="A simple script that uses the Mistral AI OCR API to get the Markdown text from a PDF or image file.",
+    epilog=example_text,
+    formatter_class=argparse.RawDescriptionHelpFormatter
+  )
+
+  parser.add_argument("input", type=Path, nargs="?", help="input PDF or image file", default=None)
+  parser.add_argument("-k", "--api-key", help="Mistral API key, can be set via the MISTRAL_API_KEY environment variable", default=None)
+  parser.add_argument("-o", "--output", type=Path, help="output directory path. If not set, a directory will be created in the current working directory using the same stem (filename without extension) as the input file", default=None)
+  parser.add_argument("-j", "--json-ocr-response", type=Path, help="path from which to load a pre-existing JSON OCR response (any input file will be ignored)", default=None)
+  parser.add_argument("-m", "--mode", type=str, choices=_mode_choices, default="FULL_NO_PAGES",
+                      help="mode of operation: either the name or numerical value of the mode. Defaults to FULL_NO_PAGES")
+  parser.add_argument("-s", "--page-separator", type=str, default="\n",
+                      help="page separator to use when writing the Markdown file. Defaults to '\\n'")
+  parser.add_argument("-n", "--no-json", action="store_false", dest="write_json",
+                      help="do not write the JSON OCR response to a file. By default, the response is written")
+  parser.add_argument("-e", "--load-dot-env", action="store_true",
+                      help="load the .env file from the current directory using python-dotenv, to retrieve the Mistral API key")
+  args = parser.parse_args()
+
+  if args.load_dot_env:
+    load_dotenv()
+
+  if args.api_key is None:
+    args.api_key = getenv("MISTRAL_API_KEY")
+    if args.api_key is None:
+      parser.error("API key is required. Set it with --api-key, via the MISTRAL_API_KEY environment variable, or load it from a .env file with -e/--load-dot-env")
+
+  try:
+    construct_from_mode(
+      document_path=args.input,
+      api_key=args.api_key,
+      output_directory=args.output,
+      json_ocr_response_path=args.json_ocr_response,
+      page_separator=_unescape(args.page_separator),
+      write_json=args.write_json,
+      mode=get_mode_from_string(args.mode)
+    ).process()
+  except FileNotFoundError as e:
+    parser.error(e)
+  except ValueError as e:
+    parser.error(e)
+
+if __name__ == "__main__":
+  main()

mistral_ai_ocr-1.0.dist-info/METADATA

@@ -0,0 +1,174 @@
+Metadata-Version: 2.1
+Name: mistral-ai-ocr
+Version: 1.0
+Description-Content-Type: text/markdown
+Requires-Dist: mistralai
+Requires-Dist: python-dotenv
+
+# Mistral AI OCR
+This is a simple script that uses the Mistral AI OCR API to extract text from a PDF or image file
+
+## Modes
+
+| Value | Name |
+|-|-|
+| 0 | FULL |
+| 1 | FULL_ALT |
+| 2 | FULL_NO_DIR |
+| 3 | FULL_NO_PAGES |
+| 4 | TEXT |
+| 5 | TEXT_NO_PAGES |
+
+Given the input file `paper.pdf`, the directory structure for each mode is shown below:
+
+### 0 - `FULL`
+
+Structure
+```
+paper
+├── full
+│   ├── image1.png
+│   ├── image2.png
+│   ├── image3.png
+│   └── paper.md
+├── page_0
+│   ├── image1.png
+│   └── paper.md
+├── page_1
+│   ├── image2.png
+│   └── paper.md
+└── page_2
+    ├── image3.png
+    └── paper.md
+```
+
+### 1 - `FULL_ALT`
+
+Structure
+```
+paper
+├── image1.png
+├── image2.png
+├── image3.png
+├── paper.md
+├── page_0
+│   ├── image1.png
+│   └── paper.md
+├── page_1
+│   ├── image2.png
+│   └── paper.md
+└── page_2
+    ├── image3.png
+    └── paper.md
+```
+
+### 2 - `FULL_NO_DIR`
+
+Structure
+```
+paper
+├── image1.png
+├── image2.png
+├── image3.png
+├── paper.md
+├── paper0.md
+├── paper1.md
+└── paper2.md
+```
+
+### 3 - `FULL_NO_PAGES` *default*
+
+Structure
+```
+paper
+├── image1.png
+├── image2.png
+├── image3.png
+└── paper.md
+```
+
+### 4 - `TEXT`
+
+Structure
+```
+paper
+├── paper.md
+├── paper0.md
+├── paper1.md
+└── paper2.md
+```
+
+### 5 - `TEXT_NO_PAGES`
+
+Structure
+```
+paper
+└── paper.md
+```
+
+By default, the JSON response from the Mistral AI OCR API is saved in the output directory. To disable JSON output, use the `-n` or `--no-json` argument. To experiment with a different **mode** without using additional API calls, reuse an existing JSON response instead of the original input file
+
+# Usage
+
+## Install the Requirements
+
+To install the necessary requirements, run the following command:
+
+```sh
+pip install mistral-ai-ocr
+```
+
+## Typical Usage
+
+```sh
+mistral-ai-ocr paper.pdf
+mistral-ai-ocr paper.pdf --api-key jrWjJE5lFketfB2sA6vvhQK2SoHQ6R39
+mistral-ai-ocr paper.pdf -o revision
+mistral-ai-ocr paper.pdf -e
+mistral-ai-ocr paper.pdf -m FULL
+mistral-ai-ocr page74.jpg -e
+mistral-ai-ocr -j paper.json
+mistral-ai-ocr -j paper.json -m TEXT_NO_PAGES -n
+```
+
+## Arguments
+
+| Argument || Description |
+|-|-|-|
+| | | input PDF or image file |
+| -k API_KEY | --api-key API_KEY | Mistral API key, can be set via the **MISTRAL_API_KEY** environment variable |
+| -o OUTPUT | --output OUTPUT | output directory path. If not set, a directory will be created in the current working directory using the same stem (filename without extension) as the input file |
+| -j JSON_OCR_RESPONSE | --json-ocr-response JSON_OCR_RESPONSE | path from which to load a pre-existing JSON OCR response (any input file will be ignored) |
+| -m MODE | --mode MODE | mode of operation: either the name or numerical value of the mode. _Defaults to FULL_NO_PAGES_ |
+| -s PAGE_SEPARATOR | --page-separator PAGE_SEPARATOR | page separator to use when writing the Markdown file. _Defaults to `\n`_ |
+| -n | --no-json | do not write the JSON OCR response to a file. By default, the response is written |
+| -e | --load-dot-env | load the .env file from the current directory using [`python-dotenv`](https://pypi.org/project/python-dotenv/), to retrieve the Mistral API key |
+
+### Mistral AI API Key
+
+To obtain an API key, you need a [Mistral AI](https://auth.mistral.ai/ui/registration) account. Then visit [https://admin.mistral.ai/organization/api-keys](https://admin.mistral.ai/organization/api-keys) and click the **Create new key** button
+
+To avoid using `-e` to load the `.env` file, you can create one at `$HOME/.mistral_ai_ocr.env` (where `$HOME` is your home directory). It will then be automatically loaded at the start of the script
+
+For example, for an user called `vavilov`, the path would look like this:
+
+* **Linux**
+  ```
+  /home/vavilov/.mistral_ai_ocr.env  
+  ```
+
+* **macOS**
+  ```
+  /Users/vavilov/.mistral_ai_ocr.env  
+  ```
+
+* **Windows**
+  ```
+  C:\Users\vavilov\.mistral_ai_ocr.env  
+  ```
+
+and the content will be something like this:
+
+```
+MISTRAL_API_KEY=jrWjJE5lFketfB2sA6vvhQK2SoHQ6R39
+```

mistral_ai_ocr-1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+mistral_ai_ocr/__init__.py,sha256=wOwicDbjQMcMWubEnPogXRyAiV6JVrJhiXZmWRkPsVw,9248
+mistral_ai_ocr/__main__.py,sha256=5Jrno0r448BT2HdNrXLi2H6RLkUqP76IV4kmB3HuJ6g,3247
+mistral_ai_ocr-1.0.dist-info/METADATA,sha256=hqq21JC0r4mFIy0vY7UYm9R_Mbcr4ebgYawnyGNiPyo,4401
+mistral_ai_ocr-1.0.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+mistral_ai_ocr-1.0.dist-info/entry_points.txt,sha256=m-ENd87vam6706-mmfzVfBq5q028TKM-7SMLUakWd-U,64
+mistral_ai_ocr-1.0.dist-info/top_level.txt,sha256=4X0WShtu4WEMtVriRP9X2Fia0ORjbAK03bRYimMvRHA,15
+mistral_ai_ocr-1.0.dist-info/RECORD,,

mistral_ai_ocr-1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.37.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

mistral_ai_ocr-1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+mistral-ai-ocr = mistral_ai_ocr.__main__:main

mistral_ai_ocr-1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+mistral_ai_ocr