ADD: ocr options inputs for all ocr engines + fix for model training options and wait submission

indico/queries/datasets.py

@@ -27,6 +27,9 @@
         EmailOptions,
         OmnipageOcrOptionsInput,
         ReadApiOcrOptionsInput,
+        ReadApiTablesV1OcrOptionsInput,
+        ReadApiTablesV2OcrOptionsInput,
+        ReadApiV2OcrOptionsInput,
     )
     from indico.typing import AnyDict, Payload
 
@@ -206,9 +209,12 @@ class CreateDataset(RequestChain["Dataset"]):
         from_local_images (bool, optional): Flag whether files are local images or not. Defaults to False.
         image_filename_col (str, optional): Image filename column. Defaults to 'filename'.
         batch_size (int, optional): Size of file batch to upload at a time. Defaults to 20.
-        ocr_engine (OcrEngine, optional): Specify an OCR engine [OMNIPAGE, READAPI, READAPI_V2, READAPI_TABLES_V1]. Defaults to None.
+        ocr_engine (OcrEngine, optional): Specify an OCR engine [OMNIPAGE, READAPI, READAPI_V2, READAPI_TABLES_V1, READAPI_TABLES_V2]. Defaults to None.
         omnipage_ocr_options (OmnipageOcrOptionsInput, optional): If using Omnipage, specify Omnipage OCR options. Defaults to None.
         read_api_ocr_options: (ReadApiOcrOptionsInput, optional): If using ReadAPI, specify ReadAPI OCR options. Defaults to None.
+        read_api_v2_ocr_options: (ReadApiV2OcrOptionsInput, optional): If using ReadAPI v2, specify ReadAPI v2 OCR options. Defaults to None.
+        read_api_tables_v1_ocr_options: (ReadApiTablesV1OcrOptionsInput, optional): If using ReadAPI tables v1, specify ReadAPI tables v1 OCR options. Defaults to None.
+        read_api_tables_v2_ocr_options: (ReadApiTablesV2OcrOptionsInput, optional): If using ReadAPI tables v2, specify ReadAPI tables v2 OCR options. Defaults to None.
         request_interval (int or float, optional): The maximum time in between retry calls when waiting. Defaults to 5 seconds.
 
     Returns:
@@ -228,6 +234,9 @@ def __init__(
         ocr_engine: "Optional[OcrEngine]" = None,
         omnipage_ocr_options: "Optional[OmnipageOcrOptionsInput]" = None,
         read_api_ocr_options: "Optional[ReadApiOcrOptionsInput]" = None,
+        read_api_v2_ocr_options: "Optional[ReadApiV2OcrOptionsInput]" = None,
+        read_api_tables_v1_ocr_options: "Optional[ReadApiTablesV1OcrOptionsInput]" = None,
+        read_api_tables_v2_ocr_options: "Optional[ReadApiTablesV2OcrOptionsInput]" = None,
         request_interval: "Union[int, float]" = 5,
         email_options: "Optional[EmailOptions]" = None,
     ):
@@ -241,11 +250,26 @@ def __init__(
         self.ocr_engine = ocr_engine
         self.omnipage_ocr_options = omnipage_ocr_options
         self.read_api_ocr_options = read_api_ocr_options
+        self.read_api_v2_ocr_options = read_api_v2_ocr_options
+        self.read_api_tables_v1_ocr_options = read_api_tables_v1_ocr_options
+        self.read_api_tables_v2_ocr_options = read_api_tables_v2_ocr_options
         self.request_interval = request_interval
         self.email_options = email_options
-        if omnipage_ocr_options is not None and read_api_ocr_options is not None:
+        if (
+            sum(
+                opt is not None
+                for opt in [
+                    omnipage_ocr_options,
+                    read_api_ocr_options,
+                    read_api_v2_ocr_options,
+                    read_api_tables_v1_ocr_options,
+                    read_api_tables_v2_ocr_options,
+                ]
+            )
+            > 1
+        ):
             raise IndicoInputError(
-                "Must supply either omnipage or readapi options but not both."
+                "Must supply only one of omnipage, readapi, readapi v2, readapi tables v1, or readapi tables v2 options."
             )
         super().__init__()
 
@@ -300,6 +324,9 @@ def requests(
             name=self.name,
             dataset_type=self.dataset_type,
             readapi_ocr_options=self.read_api_ocr_options,
+            readapi_v2_ocr_options=self.read_api_v2_ocr_options,
+            readapi_tables_v1_ocr_options=self.read_api_tables_v1_ocr_options,
+            readapi_tables_v2_ocr_options=self.read_api_tables_v2_ocr_options,
             omnipage_ocr_options=self.omnipage_ocr_options,
             ocr_engine=self.ocr_engine,
             email_options=self.email_options,
@@ -401,6 +428,9 @@ def __init__(
         ocr_engine: "Optional[OcrEngine]" = None,
         omnipage_ocr_options: "Optional[OmnipageOcrOptionsInput]" = None,
         readapi_ocr_options: "Optional[ReadApiOcrOptionsInput]" = None,
+        readapi_v2_ocr_options: "Optional[ReadApiV2OcrOptionsInput]" = None,
+        readapi_tables_v1_ocr_options: "Optional[ReadApiTablesV1OcrOptionsInput]" = None,
+        readapi_tables_v2_ocr_options: "Optional[ReadApiTablesV2OcrOptionsInput]" = None,
         email_options: "Optional[EmailOptions]" = None,
     ):
         if not dataset_type:
@@ -412,6 +442,9 @@ def __init__(
                     "ocrEngine": ocr_engine.name,
                     "omnipageOptions": omnipage_ocr_options,
                     "readapiOptions": readapi_ocr_options,
+                    "readapiV2Options": readapi_v2_ocr_options,
+                    "readapiTablesV1Options": readapi_tables_v1_ocr_options,
+                    "readapiTablesV2Options": readapi_tables_v2_ocr_options,
                 },
                 "emailOptions": email_options,
             }

indico/queries/gallery.py

@@ -4,7 +4,7 @@
 from indico.types.component_blueprint import BlueprintPage, BlueprintTags
 
 if TYPE_CHECKING:  # pragma: no cover
-    from typing import Any, Optional
+    from typing import Any, Dict, Optional, Union
 
     from indico.typing import Payload
 
@@ -53,7 +53,7 @@ class ListGallery(PagedRequestV2[BlueprintPage]):
 
     def __init__(
         self,
-        filters: "Optional[str]" = None,
+        filters: "Optional[Union[Dict[str, Any], str]]" = None,
         limit: int = 100,
         order_by: str = "name",
         desc: bool = False,

indico/queries/model_groups/model_groups.py

@@ -311,9 +311,9 @@ def __init__(
         model_training_options_json: "Optional[str]" = None
         if model_training_options:
             if isinstance(model_training_options, dict):
-                model_training_options = json.dumps(model_training_options)
+                model_training_options_json = json.dumps(model_training_options)
             else:
-                model_training_options = model_training_options
+                model_training_options_json = model_training_options
 
         predict_options_json: "Optional[str]" = None
         if predict_options:

indico/queries/submission.py

@@ -287,15 +287,22 @@ class WaitForSubmissions(RequestChain["List[Submission]"]):
         }
     """
 
-    def __init__(self, submission_ids: "List[int]", timeout: "Union[int, float]" = 60):
+    def __init__(
+        self, submission_ids: "Union[int, List[int]]", timeout: "Union[int, float]" = 60
+    ):
         if not submission_ids:
             raise IndicoInputError("Please provide submission ids")
 
         self.submission_ids = submission_ids
         self.timeout = timeout
         self.status_check = partial(ne, "PROCESSING")
+        num_submissions = (
+            1 if isinstance(self.submission_ids, int) else len(self.submission_ids)
+        )
         self.status_getter = partial(
-            ListSubmissions, submission_ids=self.submission_ids, limit=None
+            ListSubmissions,
+            submission_ids=self.submission_ids,
+            limit=num_submissions,
         )
 
     def requests(self) -> "Iterator[ListSubmissions]":

indico/types/dataset.py

@@ -75,6 +75,11 @@ class TableReadOrder(Enum):
     COLUMN = 1
 
 
+class ExcelTablesType(Enum):
+    RENDERED = 0
+    NATIVE = 1
+
+
 class OcrEngine(Enum):
     """
     Enum representing available OCR engines.
@@ -84,17 +89,36 @@ class OcrEngine(Enum):
     READAPI = 1
     READAPI_V2 = 2
     READAPI_TABLES_V1 = 3
+    READAPI_TABLES_V2 = 4
 
 
-class OmnipageOcrOptionsInput(BaseType):
+class BaseOCROptions(BaseType):
+    """
+    Base OCR options shared across engines.
+
+    Args:
+        auto_rotate(bool): Auto rotate.
+        upscale_images(bool): Scale up low resolution images.
+        spreadsheet_converter_version(int): Spreadsheet converter version.
+        languages(List[str]): List of languages to use.
+    """
+
+    auto_rotate: bool
+    upscale_images: bool
+    spreadsheet_converter_version: int
+    languages: List[str]
+
+
+class OmnipageOcrOptionsInput(BaseOCROptions):
     """
     Omnipage specific OCR options for dataset creation.
 
     Args:
-        auto_rotate(bool): auto rotate.
-        single_colum(bool): Read table as a single column.
-        upscale_images(bool): Scale up low-resolution images.
-        languages(List[OmnipageLanguageCode]): List of languages to use in ocr.
+        auto_rotate(bool): Auto rotate.
+        upscale_images(bool): Scale up low resolution images.
+        spreadsheet_converter_version(int): Spreadsheet converter version.
+        languages(List[str]): List of languages to use.
+        single_column(bool): Read table as a single column.
         cells(bool): Return table information for post-processing rules
         force_render(bool): Force rednering.
         native_layout(bool): Native layout.
@@ -103,32 +127,90 @@ class OmnipageOcrOptionsInput(BaseType):
 
     """
 
-    auto_rotate: bool
     single_column: bool
-    upscale_images: bool
-    languages: List[str]
     cells: bool
     force_render: bool
     native_layout: bool
     native_pdf: bool
     table_read_order: TableReadOrder
+    split_version: int
 
 
-class ReadApiOcrOptionsInput(BaseType):
+class ReadApiOcrOptionsInput(BaseOCROptions):
     """
     Read API OCR options.
 
     Args:
-        auto_rotate(bool): Auto rotate
-        single_column(bool): Read table as a single column.
+        auto_rotate(bool): Auto rotate.
         upscale_images(bool): Scale up low resolution images.
+        spreadsheet_converter_version(int): Spreadsheet converter version.
         languages(List[str]): List of languages to use.
+        excel_tables(bool): Enable excel tables processing.
+        excel_tables_type(ExcelTablesType): Excel tables processing type (NATIVE or RENDERED).
+        single_column(bool): Read table as a single column.
     """
 
-    auto_rotate: bool
     single_column: bool
-    upscale_images: bool
-    languages: List[str]
+    excel_tables: bool
+    excel_tables_type: ExcelTablesType
+
+
+class ReadApiV2OcrOptionsInput(ReadApiOcrOptionsInput):
+    """
+    Read API v2 OCR options.
+
+    Args:
+        auto_rotate(bool): Auto rotate.
+        upscale_images(bool): Scale up low resolution images.
+        spreadsheet_converter_version(int): Spreadsheet converter version.
+        languages(List[str]): List of languages to use.
+        excel_tables(bool): Enable excel tables processing.
+        excel_tables_type(ExcelTablesType): Excel tables processing type (NATIVE or RENDERED).
+        single_column(bool): Read table as a single column.
+    """
+
+
+class ReadApiTablesV1OcrOptionsInput(ReadApiOcrOptionsInput):
+    """
+    Read API tables v1 OCR options.
+
+    Args:
+        auto_rotate(bool): Auto rotate.
+        upscale_images(bool): Scale up low resolution images.
+        spreadsheet_converter_version(int): Spreadsheet converter version.
+        languages(List[str]): List of languages to use.
+        excel_tables(bool): Enable excel tables processing.
+        excel_tables_type(ExcelTablesType): Excel tables processing type (NATIVE or RENDERED).
+        single_column(bool): Read table as a single column.
+        table_read_order(TableReadOrder): Read table by row or column.
+    """
+
+    table_read_order: TableReadOrder
+
+
+class ReadApiTablesV2OcrOptionsInput(BaseOCROptions):
+    """
+    Read API tables v2 OCR options.
+
+    Args:
+        auto_rotate(bool): Auto rotate.
+        upscale_images(bool): Scale up low resolution images.
+        spreadsheet_converter_version(int): Spreadsheet converter version.
+        languages(List[str]): List of languages to use.
+        excel_tables(bool): Enable excel tables processing.
+        excel_tables_type(ExcelTablesType): Excel tables processing type (NATIVE or RENDERED).
+        table_read_order(TableReadOrder): Read table by row or column.
+        include_markdown(bool): Include formatted text in the output.
+        include_barcodes(bool): Recognize and extract barcodes.
+        include_key_value_pairs(bool): Recognize and extract key-value pairs.
+    """
+
+    excel_tables: bool
+    excel_tables_type: ExcelTablesType
+    table_read_order: TableReadOrder
+    include_markdown: bool
+    include_barcodes: bool
+    include_key_value_pairs: bool
 
 
 class OcrInputLanguage(BaseType):
@@ -165,3 +247,6 @@ class OcrOptionsInput:
     ocr_engine: OcrEngine
     omnipage_options: OmnipageOcrOptionsInput
     readapi_options: ReadApiOcrOptionsInput
+    readapi_v2_options: ReadApiV2OcrOptionsInput
+    readapi_tables_v1_options: ReadApiTablesV1OcrOptionsInput
+    readapi_tables_v2_options: ReadApiTablesV2OcrOptionsInput