Skip to content

Commit f80ead1

Browse files
committed
Add in layout README
1 parent 2751dcf commit f80ead1

15 files changed

+86
-29
lines changed

README.md

+72-23
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,13 @@ Surya is a document OCR toolkit that does:
44

55
- Accurate OCR in 90+ languages
66
- Line-level text detection in any language
7-
- Table and chart detection (coming soon)
7+
- Layout analysis (table, image, header, etc detection) in any language
88

99
It works on a range of documents (see [usage](#usage) and [benchmarks](#benchmarks) for more details).
1010

11-
| Detection | OCR |
12-
|:----------------------------------------------------------------:|:-----------------------------------------------------------------------:|
13-
| ![New York Times Article Detection](static/images/excerpt.png) | ![New York Times Article Recognition](static/images/excerpt_text.png) |
11+
| Detection | OCR | Layout |
12+
|:----------------------------------------------------------------:|:-----------------------------------------------------------------------:|:---------------------------------------------------------------------:|
13+
| ![New York Times Article Detection](static/images/excerpt.png) | ![New York Times Article Recognition](static/images/excerpt_text.png) | ![New York Times Article Detection](static/images/excerpt_layout.png) |
1414

1515

1616
Surya is named for the [Hindu sun god](https://en.wikipedia.org/wiki/Surya), who has universal vision.
@@ -21,19 +21,19 @@ Surya is named for the [Hindu sun god](https://en.wikipedia.org/wiki/Surya), who
2121

2222
## Examples
2323

24-
| Name | Text Detection | OCR |
25-
|------------------|:-----------------------------------:|-----------------------------------------:|
26-
| Japanese | [Image](static/images/japanese.jpg) | [Image](static/images/japanese_text.jpg) |
27-
| Chinese | [Image](static/images/chinese.jpg) | [Image](static/images/chinese_text.jpg) |
28-
| Hindi | [Image](static/images/hindi.jpg) | [Image](static/images/hindi_text.jpg) |
29-
| Arabic | [Image](static/images/arabic.jpg) | [Image](static/images/arabic_text.jpg) |
30-
| Chinese + Hindi | [Image](static/images/chi_hind.jpg) | [Image](static/images/chi_hind_text.jpg) |
31-
| Presentation | [Image](static/images/pres.png) | [Image](static/images/pres_text.jpg) |
32-
| Scientific Paper | [Image](static/images/paper.jpg) | [Image](static/images/paper_text.jpg) |
33-
| Scanned Document | [Image](static/images/scanned.png) | [Image](static/images/scanned_text.jpg) |
34-
| New York Times | [Image](static/images/nyt.jpg) | [Image](static/images/nyt_text.jpg) |
35-
| Scanned Form | [Image](static/images/funsd.png) | [Image](static/images/funsd_text.jpg) |
36-
| Textbook | [Image](static/images/textbook.jpg) | [Image](static/images/textbook_text.jpg) |
24+
| Name | Text Detection | OCR | Layout |
25+
|------------------|:-----------------------------------:|-----------------------------------------:|--------:|
26+
| Japanese | [Image](static/images/japanese.jpg) | [Image](static/images/japanese_text.jpg) | [Image](static/images/japanese_layout.jpg) |
27+
| Chinese | [Image](static/images/chinese.jpg) | [Image](static/images/chinese_text.jpg) | [Image](static/images/chinese_layout.jpg) |
28+
| Hindi | [Image](static/images/hindi.jpg) | [Image](static/images/hindi_text.jpg) | [Image](static/images/hindi_layout.jpg) |
29+
| Arabic | [Image](static/images/arabic.jpg) | [Image](static/images/arabic_text.jpg) | [Image](static/images/arabic_layout.jpg) |
30+
| Chinese + Hindi | [Image](static/images/chi_hind.jpg) | [Image](static/images/chi_hind_text.jpg) | [Image](static/images/chi_hind_layout.jpg) |
31+
| Presentation | [Image](static/images/pres.png) | [Image](static/images/pres_text.jpg) | [Image](static/images/pres_layout.jpg) |
32+
| Scientific Paper | [Image](static/images/paper.jpg) | [Image](static/images/paper_text.jpg) | [Image](static/images/paper_layout.jpg) |
33+
| Scanned Document | [Image](static/images/scanned.png) | [Image](static/images/scanned_text.jpg) | [Image](static/images/scanned_layout.jpg) |
34+
| New York Times | [Image](static/images/nyt.jpg) | [Image](static/images/nyt_text.jpg) | [Image](static/images/nyt_layout.jpg) |
35+
| Scanned Form | [Image](static/images/funsd.png) | [Image](static/images/funsd_text.jpg) | -- |
36+
| Textbook | [Image](static/images/textbook.jpg) | [Image](static/images/textbook_text.jpg) | [Image](static/images/textbook_layout.jpg) |
3737

3838
# Installation
3939

@@ -100,13 +100,13 @@ Setting the `RECOGNITION_BATCH_SIZE` env var properly will make a big difference
100100
from PIL import Image
101101
from surya.ocr import run_ocr
102102
from surya.model.detection.segformer import load_model as load_det_model, load_processor as load_det_processor
103-
from surya.model.recognition.model import load_model as load_rec_model
104-
from surya.model.recognition.processor import load_processor as load_rec_processor
103+
from surya.model.recognition.model import load_model
104+
from surya.model.recognition.processor import load_processor
105105
106106
image = Image.open(IMAGE_PATH)
107107
langs = ["en"] # Replace with your languages
108108
det_processor, det_model = load_det_processor(), load_det_model()
109-
rec_model, rec_processor = load_rec_model(), load_rec_processor()
109+
rec_model, rec_processor = load_model(), load_processor()
110110
111111
predictions = run_ocr([image], [langs], det_model, det_processor, rec_model, rec_processor)
112112
```
@@ -156,13 +156,59 @@ model, processor = load_model(), load_processor()
156156
predictions = batch_detection([image], model, processor)
157157
```
158158

159+
## Layout analysis
160+
161+
You can detect the layout of an image, pdf, or folder of images/pdfs with the following command. This will write out a json file with the detected layout.
162+
163+
```
164+
surya_layout DATA_PATH --images
165+
```
166+
167+
- `DATA_PATH` can be an image, pdf, or folder of images/pdfs
168+
- `--images` will save images of the pages and detected text lines (optional)
169+
- `--max` specifies the maximum number of pages to process if you don't want to process everything
170+
- `--results_dir` specifies the directory to save results to instead of the default
171+
172+
The `results.json` file will contain a json dictionary where the keys are the input filenames without extensions. Each value will be a list of dictionaries, one per page of the input document. Each page dictionary contains:
173+
174+
- `bboxes` - detected bounding boxes for text
175+
- `bbox` - the axis-aligned rectangle for the text line in (x1, y1, x2, y2) format. (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner.
176+
- `polygon` - the polygon for the text line in (x1, y1), (x2, y2), (x3, y3), (x4, y4) format. The points are in clockwise order from the top left.
177+
- `confidence` - the confidence of the model in the detected text (0-1). This is currently not very reliable.
178+
- `label` - the label for the bbox. One of `Caption`, `Footnote`, `Formula`, `List-item`, `Page-footer`, `Page-header`, `Picture`, `Figure`, `Section-header`, `Table`, `Text`, `Title`.
179+
- `page` - the page number in the file
180+
- `image_bbox` - the bbox for the image in (x1, y1, x2, y2) format. (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner. All line bboxes will be contained within this bbox.
181+
182+
**Performance tips**
183+
184+
Setting the `DETECTOR_BATCH_SIZE` env var properly will make a big difference when using a GPU. Each batch item will use `280MB` of VRAM, so very high batch sizes are possible. The default is a batch size `32`, which will use about 9GB of VRAM. Depending on your CPU core count, it might help, too - the default CPU batch size is `2`.
185+
186+
### From python
187+
188+
```
189+
from PIL import Image
190+
from surya.detection import batch_detection
191+
from surya.model.segformer import load_model, load_processor
192+
from surya.settings import settings
193+
194+
image = Image.open(IMAGE_PATH)
195+
model = load_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
196+
processor = load_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
197+
det_model = load_model()
198+
det_processor = load_processor()
199+
200+
# layout_predictions is a list of dicts, one per image
201+
line_predictions = batch_text_detection([image], det_model, det_processor)
202+
layout_predictions = batch_layout_detection([image], model, processor, line_predictions)
203+
```
204+
159205
# Limitations
160206

161207
- This is specialized for document OCR. It will likely not work on photos or other images.
162208
- Surya is for OCR - the goal is to recognize the text lines correctly, not sort them into reading order. Surya will attempt to sort the lines, which will work in many cases, but use something like [marker](https://github.com/VikParuchuri/marker) or other postprocessing if you need to order the text.
163209
- It is for printed text, not handwriting (though it may work on some handwriting).
164-
- The model has trained itself to ignore advertisements.
165-
- You can find language support for OCR in `surya/languages.py`. Text detection should work with any language.
210+
- The text detection model has trained itself to ignore advertisements.
211+
- You can find language support for OCR in `surya/languages.py`. Text detection and layout analysis will work with any language.
166212

167213
## Troubleshooting
168214

@@ -172,7 +218,6 @@ If OCR isn't working properly:
172218
- Preprocessing the image (binarizing, deskewing, etc) can help with very old/blurry images.
173219
- You can adjust `DETECTOR_BLANK_THRESHOLD` and `DETECTOR_TEXT_THRESHOLD` if you don't get good results. `DETECTOR_BLANK_THRESHOLD` controls the space between lines - any prediction below this number will be considered blank space. `DETECTOR_TEXT_THRESHOLD` controls how text is joined - any number above this is considered text. `DETECTOR_TEXT_THRESHOLD` should always be higher than `DETECTOR_BLANK_THRESHOLD`, and both should be in the 0-1 range. Looking at the heatmap from the debug output of the detector can tell you how to adjust these (if you see faint things that look like boxes, lower the thresholds, and if you see bboxes being joined together, raise the thresholds).
174220

175-
176221
# Manual install
177222

178223
If you want to develop surya, you can install it manually:
@@ -231,6 +276,10 @@ First calculate coverage for each bbox, then add a small penalty for double cove
231276

232277
Then we calculate precision and recall for the whole dataset.
233278

279+
## Layout analysis
280+
281+
282+
234283
## Running your own benchmarks
235284

236285
You can benchmark the performance of surya on your machine.

detect_layout.py

-2
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,6 @@ def main():
2121
parser.add_argument("--debug", action="store_true", help="Run in debug mode.", default=False)
2222
args = parser.parse_args()
2323

24-
print("Layout detection is currently in beta! There may be issues with the output.")
25-
2624
model = load_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
2725
processor = load_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
2826
det_model = load_model()

static/images/arabic_layout.jpg

302 KB
Loading

static/images/chi_hind_layout.jpg

483 KB
Loading

static/images/chinese_layout.jpg

333 KB
Loading

static/images/excerpt_layout.png

344 KB
Loading

static/images/hindi_layout.jpg

300 KB
Loading

static/images/japanese_layout.jpg

425 KB
Loading

static/images/nyt_layout.jpg

2.09 MB
Loading

static/images/paper_layout.jpg

605 KB
Loading

static/images/pres_layout.jpg

532 KB
Loading

static/images/scanned_layout.jpg

863 KB
Loading

static/images/textbook_layout.jpg

379 KB
Loading

surya/input/load.py

+11-4
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import PIL
2+
13
from surya.input.processing import open_pdf, get_page_images
24
import os
35
import filetype
@@ -50,14 +52,19 @@ def load_from_folder(folder_path, max_pages=None, start_page=None):
5052
images = []
5153
names = []
5254
for path in image_paths:
53-
if filetype.guess(path).extension == "pdf":
55+
extension = filetype.guess(path)
56+
if extension and extension.extension == "pdf":
5457
image, name = load_pdf(path, max_pages, start_page)
5558
images.extend(image)
5659
names.extend(name)
5760
else:
58-
image, name = load_image(path)
59-
images.extend(image)
60-
names.extend(name)
61+
try:
62+
image, name = load_image(path)
63+
images.extend(image)
64+
names.extend(name)
65+
except PIL.UnidentifiedImageError:
66+
print(f"Could not load image {path}")
67+
continue
6168
return images, names
6269

6370

surya/layout.py

+3
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,9 @@ def get_regions_from_detection_result(detection_result: TextDetectionResult, hea
5858

5959
new_boxes = []
6060
for bbox_idx, bbox in enumerate(detected_boxes):
61+
if bbox.label == "Picture" and bbox.area < 200: # Remove very small figures
62+
continue
63+
6164
if bbox_idx not in box_lines and bbox.label not in ["Picture", "Formula"]:
6265
continue
6366

0 commit comments

Comments
 (0)