From 0e9475fc5647c461a3fba79396818b69fe895736 Mon Sep 17 00:00:00 2001 From: nastyboget Date: Mon, 22 Jun 2026 15:50:24 +0300 Subject: [PATCH 1/3] TALIE-1535: add orientation classification --- .github/workflows/test_on_push.yaml | 2 +- CHANGELOG.md | 4 + VERSION | 2 +- .../orientation_classification/__init__.py | 3 + .../orientation_classification/model.py | 23 +++++ .../orientation_classifier.py | 84 +++++++++++++++++++ pyproject.toml | 20 +++-- 7 files changed, 130 insertions(+), 8 deletions(-) create mode 100644 dedocutils/preprocessing/orientation_classification/__init__.py create mode 100644 dedocutils/preprocessing/orientation_classification/model.py create mode 100644 dedocutils/preprocessing/orientation_classification/orientation_classifier.py diff --git a/.github/workflows/test_on_push.yaml b/.github/workflows/test_on_push.yaml index 92f7100..8b57867 100644 --- a/.github/workflows/test_on_push.yaml +++ b/.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] diff --git a/CHANGELOG.md b/CHANGELOG.md index cdd0f83..adeec4a 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/VERSION b/VERSION index 4209dba..ed63cdf 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.3.8 \ No newline at end of file +0.3.9 \ No newline at end of file diff --git a/dedocutils/preprocessing/orientation_classification/__init__.py b/dedocutils/preprocessing/orientation_classification/__init__.py new file mode 100644 index 0000000..3700ebc --- /dev/null +++ b/dedocutils/preprocessing/orientation_classification/__init__.py @@ -0,0 +1,3 @@ +from .orientation_classifier import OrientationClassifier + +__all__ = ["OrientationClassifier"] diff --git a/dedocutils/preprocessing/orientation_classification/model.py b/dedocutils/preprocessing/orientation_classification/model.py new file mode 100644 index 0000000..90681cd --- /dev/null +++ b/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): + """ + Class detects EfficientNet B0 model + """ + def __init__(self, model_path: Optional[str], num_classes: int = 6) -> None: + """ + first 2 classes mean columns number + last 4 classes mean orientation + """ + 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 diff --git a/dedocutils/preprocessing/orientation_classification/orientation_classifier.py b/dedocutils/preprocessing/orientation_classification/orientation_classifier.py new file mode 100644 index 0000000..794bd06 --- /dev/null +++ b/dedocutils/preprocessing/orientation_classification/orientation_classifier.py @@ -0,0 +1,84 @@ +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. This class set device, + preprocessing (transform) input data, weights of model + """ + def __init__(self, checkpoint_path: str) -> None: + if torch.cuda.is_available(): + self.device = torch.device("cuda:0") + self.location = lambda storage, loc: storage.cuda() + else: + self.device = torch.device("cpu") + self.location = "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 + + @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().to(self.device) + + for i in range(0, len(tensor_images), batch_size): + batch = tensor_images[i:i + batch_size] + outputs = self.net(batch) + + # 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 diff --git a/pyproject.toml b/pyproject.toml index a2e1a8c..05ae475 100644 --- a/pyproject.toml +++ b/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" From 877ae0e3d845886ae82e8c05df0f2d104fddf5a6 Mon Sep 17 00:00:00 2001 From: nastyboget Date: Tue, 23 Jun 2026 13:51:33 +0300 Subject: [PATCH 2/3] TALIE-1535: fixes for gpu --- .../orientation_classifier.py | 23 ++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/dedocutils/preprocessing/orientation_classification/orientation_classifier.py b/dedocutils/preprocessing/orientation_classification/orientation_classifier.py index 794bd06..ad5e09b 100644 --- a/dedocutils/preprocessing/orientation_classification/orientation_classifier.py +++ b/dedocutils/preprocessing/orientation_classification/orientation_classifier.py @@ -16,13 +16,11 @@ class OrientationClassifier: Class Classifier for work with Orientation Network. This class set device, preprocessing (transform) input data, weights of model """ - def __init__(self, checkpoint_path: str) -> None: - if torch.cuda.is_available(): - self.device = torch.device("cuda:0") - self.location = lambda storage, loc: storage.cuda() + 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.device = torch.device("cpu") - self.location = "cpu" + self.to(torch.device("cpu")) logger.warning(f"Classifier is set to device {self.device}") @@ -36,6 +34,15 @@ def __init__(self, checkpoint_path: str) -> None: 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: @@ -58,11 +65,11 @@ def predict(self, images: list[Image], batch_size: int = 32) -> list[int]: all_orientation_predicted = [] with torch.no_grad(): - tensor_images = torch.stack(self.transform(images)).float().to(self.device) + 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) + outputs = self.net(batch.to(self.device)) # first 2 classes mean columns number # last 4 classes mean orientation From b7fdb3dfcc5e0cb35a384daa64f48c62cd397dae Mon Sep 17 00:00:00 2001 From: nastyboget Date: Mon, 29 Jun 2026 13:32:48 +0300 Subject: [PATCH 3/3] TALIE-1535: review fixes --- .../preprocessing/orientation_classification/model.py | 6 +++--- .../orientation_classification/orientation_classifier.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/dedocutils/preprocessing/orientation_classification/model.py b/dedocutils/preprocessing/orientation_classification/model.py index 90681cd..18664b3 100644 --- a/dedocutils/preprocessing/orientation_classification/model.py +++ b/dedocutils/preprocessing/orientation_classification/model.py @@ -7,12 +7,12 @@ class ClassificationModelTorch(nn.Module): """ - Class detects EfficientNet B0 model + Wrapper for EfficientNet B0 model. """ def __init__(self, model_path: Optional[str], num_classes: int = 6) -> None: """ - first 2 classes mean columns number - last 4 classes mean orientation + 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) diff --git a/dedocutils/preprocessing/orientation_classification/orientation_classifier.py b/dedocutils/preprocessing/orientation_classification/orientation_classifier.py index ad5e09b..f990fc5 100644 --- a/dedocutils/preprocessing/orientation_classification/orientation_classifier.py +++ b/dedocutils/preprocessing/orientation_classification/orientation_classifier.py @@ -13,8 +13,8 @@ class OrientationClassifier: """ - Class Classifier for work with Orientation Network. This class set device, - preprocessing (transform) input data, weights of model + 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():