TALIE-1535: add orientation classification

.github/workflows/test_on_push.yaml

@@ -22,7 +22,7 @@ jobs:
     - name: Install dependencies
       run: |
         python3 -m pip install --upgrade pip
-        pip3 install .[dev,torch,doctr]
+        pip3 install .[dev,torch,doctr,line-segmentation,preprocessing,tesseract]
     - name: Run lint
       run: |
         pip3 install .[lint]

CHANGELOG.md

@@ -1,6 +1,10 @@
 Changelog
 =========
 
+v0.3.9 (2026-06-19)
+-------------------
+* Add `OrientationClassifier` class
+
 v0.3.8 (2024-09-05)
 -------------------
 * Add `shift` method to `BBox` class

VERSION

@@ -1 +1 @@
-0.3.8
+0.3.9

dedocutils/preprocessing/orientation_classification/__init__.py

@@ -0,0 +1,3 @@
+from .orientation_classifier import OrientationClassifier
+
+__all__ = ["OrientationClassifier"]

dedocutils/preprocessing/orientation_classification/model.py

@@ -0,0 +1,23 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from torchvision import models
+
+
+class ClassificationModelTorch(nn.Module):
+    """
+    Wrapper for EfficientNet B0 model.
+    """
+    def __init__(self, model_path: Optional[str], num_classes: int = 6) -> None:
+        """
+        First 2 classes are the number of columns on the page [1 column, 2 columns].
+        Last 4 classes are the page orientation in degrees [0, 90, 180, 270].
+        """
+        super(ClassificationModelTorch, self).__init__()
+        self.efficientnet_b0 = models.efficientnet_b0(pretrained=model_path is None)
+        self.efficientnet_b0.classifier[1] = nn.Linear(in_features=1280, out_features=num_classes)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.efficientnet_b0(x)
+        return out

dedocutils/preprocessing/orientation_classification/orientation_classifier.py

@@ -0,0 +1,91 @@
+import logging
+import warnings
+
+import torch
+from PIL import Image
+from torchvision.transforms import v2
+from torchvision.transforms.functional import resize
+
+from dedocutils.preprocessing.orientation_classification.model import ClassificationModelTorch
+
+logger = logging.getLogger()
+
+
+class OrientationClassifier:
+    """
+    Class Classifier for work with Orientation Network for detecting document page orientation.
+    This class set device, preprocessing (transform) input data, weights of model.
+    """
+    def __init__(self, checkpoint_path: str, use_gpu: bool = True) -> None:
+        if use_gpu and torch.cuda.is_available():
+            self.to(torch.device("cuda"))
+        else:
+            self.to(torch.device("cpu"))
+
+        logger.warning(f"Classifier is set to device {self.device}")
+
+        self.transform = v2.Compose([
+            v2.Lambda(_image_resize),
+            v2.ToTensor(),
+            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        ])
+
+        self.checkpoint_path = checkpoint_path
+        self.classes = [0, 90, 180, 270]
+        self._net = None
+
+    def to(self, device: torch.device) -> None:
+        self.device = device
+        if device == torch.device("cpu"):
+            self.location = "cpu"
+        else:
+            self.location = lambda storage, loc: storage.cuda()
+
+        self._net = None
+
+    @property
+    def net(self) -> ClassificationModelTorch:
+        if self._net:
+            return self._net
+
+        self._net = ClassificationModelTorch(self.checkpoint_path)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            self._net.load_state_dict(torch.load(self.checkpoint_path, map_location=self.location))
+            logger.info(f"Weights were loaded from {self.checkpoint_path}")
+
+        self._net.to(self.device)
+        self._net.eval()
+        return self._net
+
+    def predict(self, images: list[Image], batch_size: int = 32) -> list[int]:
+        """
+        Predict class orientation of input image
+        """
+        all_orientation_predicted = []
+
+        with torch.no_grad():
+            tensor_images = torch.stack(self.transform(images)).float()
+
+            for i in range(0, len(tensor_images), batch_size):
+                batch = tensor_images[i:i + batch_size]
+                outputs = self.net(batch.to(self.device))
+
+                # first 2 classes mean columns number
+                # last 4 classes mean orientation
+                _, orientation_out = outputs[:, :2], outputs[:, 2:]
+                _, orientation_predicted = torch.max(orientation_out, 1)
+                all_orientation_predicted.append(orientation_predicted)
+
+        all_orientation_predicted = torch.cat(all_orientation_predicted, dim=0)
+        predicted_angles = [self.classes[int(predicted_angle)] for predicted_angle in all_orientation_predicted]
+        logger.info(f"Predicted orientation: {predicted_angles}")
+        return predicted_angles
+
+
+def _image_resize(image: Image) -> Image:
+    max_dim = max(image.size)
+    image1 = resize(image, size=[round(image.size[1] / max_dim * 1200), round(image.size[0] / max_dim * 1200)])
+    white_image = Image.new(size=(1200, 1200), color=(255, 255, 255), mode="RGB")
+    white_image.paste(image1)
+    return white_image

pyproject.toml

@@ -18,6 +18,7 @@ classifiers=[
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
         "Programming Language :: Python :: 3 :: Only",
 ]
 description = "Utils for automatic document images processing"
@@ -27,18 +28,25 @@ license = {file = "LICENSE"}
 dynamic = ["version"]
 requires-python = ">=3.6"
 dependencies = [
-    "numpy>=1.22.0,<2.0",
-    "opencv-python>=4.5.5.64,<=4.6.0.66",
-    "pytesseract>=0.3",
-    "scikit_learn>=1.0.2,<=1.3.1"
+    "numpy>=1.22.0,<3.0"
 ]
 
 [project.optional-dependencies]
+line-segmentation = [
+    "scikit_learn>=1.0.2"
+]
+preprocessing = [
+    "opencv-python>=4.5.5.64"
+]
+tesseract = [
+    "pytesseract>=0.3"
+]
 torch = [
-    "torch~=1.11.0",
-    "torchvision~=0.12.0"
+    "torch>=1.11.0,<3.0",
+    "torchvision>=0.12.0,<1.0"
 ]
 doctr = [
+    "numpy<2",
     "pyclipper==1.3.0.post4",
     "shapely==2.0.1",
     "tqdm>=4"