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..18664b3 --- /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): + """ + 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 diff --git a/dedocutils/preprocessing/orientation_classification/orientation_classifier.py b/dedocutils/preprocessing/orientation_classification/orientation_classifier.py new file mode 100644 index 0000000..f990fc5 --- /dev/null +++ b/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 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"