Add (Recognition|Segmentation)BaseModel classes

Author: mittagessen

kraken/models/__init__.py

@@ -1,4 +1,4 @@
-from .base import BaseModel  # NOQA
+from .base import *  # NOQA
 from .writers import *  # NOQA
 from .loaders import *  # NOQA
 from .utils import *  # NOQA

kraken/models/base.py

@@ -7,9 +7,12 @@
 import logging
 
 from abc import ABC, abstractmethod
+from collections.abc import Generator
 from typing import Union, Literal, NewType, TYPE_CHECKING
 
-__all__ = ['BaseModel']
+__all__ = ['BaseModel',
+           'RecognitionBaseModel',
+           'SegmentationBaseModel']
 
 if TYPE_CHECKING:
     from kraken.configs import Config
@@ -87,3 +90,24 @@ def prepare_for_inference(self, config: 'Config'):
         Prepares the model for inference.
         """
         pass
+
+
+class SegmentationBaseModel(BaseModel):
+    """
+    Base model metaclass for layout analysis models.
+    """
+    @abstractmethod
+    def predict(self, im: 'Image.Image') -> 'Segmentation':
+        """
+        Computes a segmentation on an input image.
+        """
+        pass
+
+
+class RecognitionBaseModel(BaseModel):
+    """
+    Base model metaclass for text recognition models.
+    """
+    @abstractmethod
+    def predict(self, im: 'Image.Image', segmentation: 'Segmentation') -> Generator['ocr_record', None, None]:
+        pass