manuscript-ocr 0.1.0__tar.gz

manuscript_ocr-0.1.0/MANIFEST.in

@@ -0,0 +1,2 @@
+include README.md LICENSE requirements.txt
+recursive-include src *

manuscript_ocr-0.1.0/PKG-INFO

@@ -0,0 +1,73 @@
+Metadata-Version: 2.4
+Name: manuscript-ocr
+Version: 0.1.0
+Summary: manuscript-ocr
+Home-page: 
+Author: 
+Author-email: 
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: torch==2.0.1+cu118
+Requires-Dist: torchvision==0.15.2+cu118
+Requires-Dist: torchaudio==2.0.2
+Requires-Dist: numpy<2
+Requires-Dist: opencv-python
+Requires-Dist: Pillow
+Requires-Dist: shapely
+Requires-Dist: numba
+Requires-Dist: tensorboard
+Requires-Dist: gdown
+Requires-Dist: pydantic
+Requires-Dist: scikit-image
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+## Installation
+
+```bash
+pip install manuscript-ocr
+````
+
+## Usage Example
+
+```python
+from PIL import Image
+from manuscript.detectors import EASTInfer
+
+# Инициализация
+det = EASTInfer(score_thresh=0.9)
+
+# Инфер с визуализацией
+page, vis_image = det.infer(r"example\ocr_example_image.jpg", vis=True)
+
+print(page)
+
+# Покажет картинку с наложенными боксами
+Image.fromarray(vis_image).show()
+
+# Или сохранить результат на диск:
+Image.fromarray(vis_image).save(r"example\ocr_example_image_infer.png")
+```
+
+### Результат
+
+Текстовые блоки будут выведены в консоль, например:
+
+```
+Page(blocks=[Block(words=[Word(polygon=[(874.1005, 909.1005), (966.8995, 909.1005), (966.8995, 956.8995), (874.1005, 956.8995)]),
+                          Word(polygon=[(849.1234, 810.5678), … ])])])
+```
+
+А визуализация сохранится в файл `example/ocr_example_image_infer.png`:
+
+![OCR Inference Result](example/ocr_example_image_infer.png)

manuscript_ocr-0.1.0/README.md

@@ -0,0 +1,39 @@
+## Installation
+
+```bash
+pip install manuscript-ocr
+````
+
+## Usage Example
+
+```python
+from PIL import Image
+from manuscript.detectors import EASTInfer
+
+# Инициализация
+det = EASTInfer(score_thresh=0.9)
+
+# Инфер с визуализацией
+page, vis_image = det.infer(r"example\ocr_example_image.jpg", vis=True)
+
+print(page)
+
+# Покажет картинку с наложенными боксами
+Image.fromarray(vis_image).show()
+
+# Или сохранить результат на диск:
+Image.fromarray(vis_image).save(r"example\ocr_example_image_infer.png")
+```
+
+### Результат
+
+Текстовые блоки будут выведены в консоль, например:
+
+```
+Page(blocks=[Block(words=[Word(polygon=[(874.1005, 909.1005), (966.8995, 909.1005), (966.8995, 956.8995), (874.1005, 956.8995)]),
+                          Word(polygon=[(849.1234, 810.5678), … ])])])
+```
+
+А визуализация сохранится в файл `example/ocr_example_image_infer.png`:
+
+![OCR Inference Result](example/ocr_example_image_infer.png)

manuscript_ocr-0.1.0/pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"

manuscript_ocr-0.1.0/requirements.txt

@@ -0,0 +1,13 @@
+torch==2.0.1+cu118
+torchvision==0.15.2+cu118
+torchaudio==2.0.2
+
+numpy<2
+opencv-python
+Pillow
+shapely
+numba
+tensorboard
+gdown
+pydantic
+scikit-image

manuscript_ocr-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

manuscript_ocr-0.1.0/setup.py

@@ -0,0 +1,35 @@
+import os
+from setuptools import setup, find_packages
+
+def parse_requirements(fname="requirements.txt"):
+    here = os.path.dirname(__file__)
+    with open(os.path.join(here, fname)) as f:
+        return [ln.strip() for ln in f if ln.strip() and not ln.startswith("#")]
+
+setup(
+    name="manuscript-ocr",
+    version="0.1.0",
+    description="manuscript-ocr",
+    long_description=open("README.md", encoding="utf-8").read(),
+    long_description_content_type="text/markdown",
+
+    author="",
+    author_email="",
+
+    url="",
+    license="MIT",
+
+    packages=find_packages(where="src"),
+    package_dir={"": "src"},
+
+    python_requires=">=3.8",
+    install_requires=parse_requirements(),
+
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+    include_package_data=True,
+)

manuscript_ocr-0.1.0/src/example.py

@@ -0,0 +1,13 @@
+from PIL import Image
+from manuscript.detectors import EASTInfer
+
+# инициализация
+det = EASTInfer(score_thresh=0.9)
+
+# инфер
+page, image = det.infer(r"example\ocr_example_image.jpg", vis=True)
+print(page)
+
+pil_img = Image.fromarray(image)
+
+pil_img.show()

manuscript_ocr-0.1.0/src/manuscript/detectors/__init__.py

@@ -0,0 +1,3 @@
+from .east import EASTInfer
+
+__all__ = ["EASTInfer"]

manuscript_ocr-0.1.0/src/manuscript/detectors/east/__init__.py

@@ -0,0 +1,112 @@
+import cv2
+import numpy as np
+import torch
+from torchvision import transforms
+from pathlib import Path
+from typing import Union, Optional, List, Tuple
+
+from .east import TextDetectionFCN
+from .utils import decode_boxes_from_maps, draw_boxes
+from ..types import Word, Block, Page
+import os
+
+import gdown
+
+class EASTInfer:
+    def __init__(
+        self,
+        weights_path: Optional[Union[str, Path]] = None,
+        device: Optional[str] = None,
+        target_size: int = 1024,
+        score_geo_scale: float = 0.25,
+        shrink_ratio: float = 0.5,
+        score_thresh: float = 0.9,
+        iou_threshold: float = 0.2,
+    ):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+        if weights_path is None:
+            url = (
+                "https://github.com/konstantinkozhin/manuscript-ocr"
+                "/releases/download/v0.1.0/east_quad_14_05.pth"
+            )
+            out = os.path.expanduser("~/.east_weights.pth")
+            if not os.path.exists(out):
+                print(f"Downloading EAST weights from {url} …")
+                gdown.download(url, out, quiet=False)
+            weights_path = out
+        print(weights_path)
+        # Загружаем модель с весами
+        self.model = TextDetectionFCN(
+            pretrained_backbone=False,
+            pretrained_model_path=str(weights_path),
+        ).to(self.device)
+        self.model.eval()
+
+        self.target_size = target_size
+        self.score_geo_scale = score_geo_scale
+        self.shrink_ratio = shrink_ratio
+        self.score_thresh = score_thresh
+        self.iou_threshold = iou_threshold
+
+        self.tf = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=(0.5, 0.5, 0.5),
+                                 std=(0.5, 0.5, 0.5)),
+        ])
+
+    def infer(
+        self,
+        img_or_path: Union[str, Path, np.ndarray],
+        vis: bool = False
+    ) -> Union[Page, Tuple[Page, np.ndarray]]:
+        """
+        :param img_or_path: путь или RGB ndarray
+        :param vis: если True, возвращает также изображение с боксами
+        :return: Page или (Page, vis_image)
+        """
+        # 1) Read & RGB
+        if isinstance(img_or_path, (str, Path)):
+            img = cv2.imread(str(img_or_path))
+            if img is None:
+                raise FileNotFoundError(f"Cannot read image: {img_or_path}")
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        elif isinstance(img_or_path, np.ndarray):
+            img = img_or_path
+        else:
+            raise TypeError(f"Unsupported type {type(img_or_path)}")
+
+        # 2) Resize + ToTensor + Normalize
+        resized = cv2.resize(img, (self.target_size, self.target_size))
+        img_t = self.tf(resized).to(self.device)
+
+        # 3) Forward
+        with torch.no_grad():
+            out = self.model(img_t.unsqueeze(0))
+
+        # 4) Extract maps
+        score_map = out["score"][0].cpu().numpy().squeeze(0)
+        geo_map = out["geometry"][0].cpu().numpy().transpose(1, 2, 0)
+
+        # 5) Decode quads
+        quads9 = decode_boxes_from_maps(
+            score_map=score_map,
+            geo_map=geo_map,
+            score_thresh=self.score_thresh,
+            scale=1.0 / self.score_geo_scale,
+            iou_threshold=self.iou_threshold,
+            expand_ratio=self.shrink_ratio,
+        )
+
+        # 6) Build Page
+        words: List[Word] = []
+        for quad in quads9:
+            pts = quad[:8].reshape(4, 2).tolist()
+            words.append(Word(polygon=pts))
+        page = Page(blocks=[Block(words=words)])
+
+        # 7) Optional visualization
+        if vis:
+            vis_img = draw_boxes(resized, quads9)
+            return page, vis_img
+        return page

manuscript_ocr-0.1.0/src/manuscript/detectors/east/dataset.py

@@ -0,0 +1,154 @@
+import os
+import json
+import cv2
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import math
+from shapely.geometry import Polygon
+import torchvision.transforms as transforms
+import skimage.draw
+from .utils import quad_to_rbox
+
+
+def order_vertices_clockwise(poly):
+    poly = np.array(poly).reshape(-1, 2)
+    s = poly.sum(axis=1)
+    diff = np.diff(poly, axis=1).flatten()
+    tl = poly[np.argmin(s)]
+    br = poly[np.argmax(s)]
+    tr = poly[np.argmin(diff)]
+    bl = poly[np.argmax(diff)]
+    return np.array([tl, tr, br, bl], dtype=np.float32)
+
+
+def shrink_poly(poly, shrink_ratio=0.3):
+    poly = np.array(poly, dtype=np.float32).reshape(-1, 2)
+    N = poly.shape[0]
+    if N != 4:
+        raise ValueError("Expected quadrilateral with 4 vertices")
+    # signed area sign
+    area = 0.0
+    for i in range(N):
+        x1, y1 = poly[i]
+        x2, y2 = poly[(i + 1) % N]
+        area += x1 * y2 - x2 * y1
+    sign = 1.0 if area > 0 else -1.0
+    new_poly = np.zeros_like(poly)
+    for i in range(N):
+        p_prev = poly[(i - 1) % N]
+        p_curr = poly[i]
+        p_next = poly[(i + 1) % N]
+        edge1 = p_curr - p_prev
+        len1 = np.linalg.norm(edge1)
+        n1 = sign * np.array([edge1[1], -edge1[0]]) / (len1 + 1e-6)
+        edge2 = p_next - p_curr
+        len2 = np.linalg.norm(edge2)
+        n2 = sign * np.array([edge2[1], -edge2[0]]) / (len2 + 1e-6)
+        n_avg = n1 + n2
+        norm_n = np.linalg.norm(n_avg)
+        if norm_n > 0:
+            n_avg /= norm_n
+        offset = shrink_ratio * min(len1, len2)
+        new_poly[i] = p_curr - offset * n_avg
+    return new_poly.astype(np.float32)
+
+
+class EASTDataset(Dataset):
+    def __init__(
+        self,
+        images_folder,
+        coco_annotation_file,
+        target_size=512,
+        score_geo_scale=0.25,
+        transform=None,
+    ):
+        self.images_folder = images_folder
+        self.target_size = target_size
+        self.score_geo_scale = score_geo_scale
+        # transform pipeline
+        if transform is None:
+            self.transform = transforms.Compose(
+                [
+                    transforms.ToPILImage(),
+                    transforms.ColorJitter(0.5, 0.5, 0.5, 0.25),
+                    transforms.ToTensor(),
+                    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+                ]
+            )
+        else:
+            self.transform = transform
+        # load COCO annotations
+        with open(coco_annotation_file, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        self.images_info = {img["id"]: img for img in data["images"]}
+        self.image_ids = list(self.images_info.keys())
+        self.annots = {}
+        for ann in data["annotations"]:
+            self.annots.setdefault(ann["image_id"], []).append(ann)
+
+    def __len__(self):
+        return len(self.image_ids)
+
+    def __getitem__(self, idx):
+        image_id = self.image_ids[idx]
+        info = self.images_info[image_id]
+        path = os.path.join(self.images_folder, info["file_name"])
+        img = cv2.imread(path)
+        if img is None:
+            raise FileNotFoundError(f"Image not found: {path}")
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        # resize
+        img_resized = cv2.resize(img, (self.target_size, self.target_size))
+        # scale annotations
+        anns = self.annots.get(image_id, [])
+        quads = []
+        for ann in anns:
+            if "segmentation" not in ann:
+                continue
+            seg = ann["segmentation"]
+            pts = np.array(seg, dtype=np.float32).reshape(-1, 2)
+            # вот тут — минимальный прямоугольник
+            rect = cv2.minAreaRect(pts)
+            box = cv2.boxPoints(rect)
+            quad = order_vertices_clockwise(box)
+            # масштабируем quad под resized
+            quad[:, 0] *= self.target_size / info["width"]
+            quad[:, 1] *= self.target_size / info["height"]
+            quads.append(quad)
+        # generate maps
+        score_map, geo_map = self.compute_quad_maps(quads)
+        rboxes = np.stack([quad_to_rbox(q.flatten()) for q in quads], axis=0).astype(
+            np.float32
+        )
+        # transform image
+        img_tensor = self.transform(img_resized)
+        target = {
+            "score_map": torch.tensor(score_map).unsqueeze(0),
+            "geo_map": torch.tensor(geo_map),
+            "rboxes": torch.tensor(rboxes),
+        }
+        return img_tensor, target
+
+    def compute_quad_maps(self, quads):
+        """
+        quads: list of (4,2) arrays
+        returns score_map (H',W') and geo_map (8,H',W')
+        """
+        out_h = int(self.target_size * self.score_geo_scale)
+        out_w = int(self.target_size * self.score_geo_scale)
+        score_map = np.zeros((out_h, out_w), dtype=np.float32)
+        geo_map = np.zeros((8, out_h, out_w), dtype=np.float32)
+        for quad in quads:
+            # shrink & order
+            shrunk = shrink_poly(order_vertices_clockwise(quad), shrink_ratio=0.3)
+            # scale to map resolution
+            coords = shrunk * self.score_geo_scale
+            rr, cc = skimage.draw.polygon(
+                coords[:, 1], coords[:, 0], shape=(out_h, out_w)
+            )
+            score_map[rr, cc] = 1
+            for i, (vx, vy) in enumerate(coords):
+                geo_map[2 * i, rr, cc] = vx - cc
+                geo_map[2 * i + 1, rr, cc] = vy - rr
+        return score_map, geo_map

manuscript_ocr-0.1.0/src/manuscript/detectors/east/east.py

@@ -0,0 +1,115 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import resnet50, ResNet50_Weights
+from torchvision.models.feature_extraction import create_feature_extractor
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, in_channels, mid_channels, out_channels):
+        super(DecoderBlock, self).__init__()
+        self.conv1x1 = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=1),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+        )
+        self.conv3x3 = nn.Sequential(
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        x = self.conv1x1(x)
+        x = self.conv3x3(x)
+        return x
+
+
+class ResNetFeatureExtractor(nn.Module):
+    def __init__(self, pretrained=True, freeze_first=False):
+        super(ResNetFeatureExtractor, self).__init__()
+        self.model = resnet50(weights=ResNet50_Weights.DEFAULT if pretrained else None)
+        if freeze_first:
+            for name, param in self.model.named_parameters():
+                if name.startswith(("conv1", "bn1", "layer1")):
+                    param.requires_grad = False
+        self.extractor = create_feature_extractor(
+            self.model,
+            return_nodes={
+                "layer1": "res1",  # stride 4, 256 channels
+                "layer2": "res2",  # stride 8, 512 channels
+                "layer3": "res3",  # stride 16,1024 channels
+                "layer4": "res4",  # stride 32,2048 channels
+            },
+        )
+
+    def forward(self, x):
+        return self.extractor(x)
+
+
+class FeatureMergingBranchResNet(nn.Module):
+    def __init__(self):
+        super(FeatureMergingBranchResNet, self).__init__()
+        self.block1 = DecoderBlock(in_channels=2048, mid_channels=512, out_channels=512)
+        self.block2 = DecoderBlock(
+            in_channels=512 + 1024, mid_channels=256, out_channels=256
+        )
+        self.block3 = DecoderBlock(
+            in_channels=256 + 512, mid_channels=128, out_channels=128
+        )
+        self.block4 = DecoderBlock(
+            in_channels=128 + 256, mid_channels=64, out_channels=32
+        )
+
+    def forward(self, feats):
+        f1 = feats["res1"]
+        f2 = feats["res2"]
+        f3 = feats["res3"]
+        f4 = feats["res4"]
+        h4 = self.block1(f4)
+        h4_up = F.interpolate(h4, scale_factor=2, mode="bilinear", align_corners=False)
+        h3 = self.block2(torch.cat([h4_up, f3], dim=1))
+        h3_up = F.interpolate(h3, scale_factor=2, mode="bilinear", align_corners=False)
+        h2 = self.block3(torch.cat([h3_up, f2], dim=1))
+        h2_up = F.interpolate(h2, scale_factor=2, mode="bilinear", align_corners=False)
+        h1 = self.block4(torch.cat([h2_up, f1], dim=1))
+        return h1
+
+
+class OutputHead(nn.Module):
+    def __init__(self):
+        super(OutputHead, self).__init__()
+        self.score_map = nn.Conv2d(32, 1, kernel_size=1)
+        self.geo_map = nn.Conv2d(32, 8, kernel_size=1)
+
+    def forward(self, x):
+        score = torch.sigmoid(self.score_map(x))
+        geometry = self.geo_map(x)
+        return score, geometry
+
+
+class TextDetectionFCN(nn.Module):
+    def __init__(
+        self, pretrained_backbone=True, freeze_first=False, pretrained_model_path=None
+    ):
+        super(TextDetectionFCN, self).__init__()
+        # ResNet50 backbone
+        self.backbone = ResNetFeatureExtractor(
+            pretrained=pretrained_backbone, freeze_first=freeze_first
+        )
+        self.decoder = FeatureMergingBranchResNet()
+        self.output_head = OutputHead()
+        # scales for maps
+        self.score_scale = 0.25
+        self.geo_scale = 0.25
+        # load optional pretrained model weights
+        if pretrained_model_path:
+            state = torch.load(pretrained_model_path, map_location="cpu")
+            self.load_state_dict(state, strict=False)
+            print(f"Loaded pretrained model from {pretrained_model_path}")
+
+    def forward(self, x):
+        feats = self.backbone(x)
+        merged = self.decoder(feats)
+        score, geometry = self.output_head(merged)
+        return {"score": score, "geometry": geometry}