Fix bugs with RGB

Author: VikParuchuri

benchmark/detection.py

@@ -7,7 +7,7 @@
 from surya.benchmark.metrics import precision_recall
 from surya.benchmark.tesseract import tesseract_parallel
 from surya.model.detection.segformer import load_model, load_processor
-from surya.input.processing import open_pdf, get_page_images
+from surya.input.processing import open_pdf, get_page_images, convert_if_not_rgb
 from surya.detection import batch_text_detection
 from surya.postprocessing.heatmap import draw_polys_on_image
 from surya.postprocessing.util import rescale_bbox
@@ -47,7 +47,7 @@ def main():
         # These have already been shuffled randomly, so sampling from the start is fine
         dataset = datasets.load_dataset(settings.DETECTOR_BENCH_DATASET_NAME, split=f"train[:{args.max}]")
         images = list(dataset["image"])
-        images = [i.convert("RGB") for i in images]
+        images = convert_if_not_rgb(images)
         correct_boxes = []
         for i, boxes in enumerate(dataset["bboxes"]):
             img_size = images[i].size

benchmark/layout.py

@@ -6,7 +6,7 @@
 from surya.benchmark.metrics import precision_recall
 from surya.detection import batch_text_detection
 from surya.model.detection.segformer import load_model, load_processor
-from surya.input.processing import open_pdf, get_page_images
+from surya.input.processing import open_pdf, get_page_images, convert_if_not_rgb
 from surya.layout import batch_layout_detection
 from surya.postprocessing.heatmap import draw_polys_on_image, draw_bboxes_on_image
 from surya.postprocessing.util import rescale_bbox
@@ -33,7 +33,7 @@ def main():
     # These have already been shuffled randomly, so sampling from the start is fine
     dataset = datasets.load_dataset(settings.LAYOUT_BENCH_DATASET_NAME, split=f"train[:{args.max}]")
     images = list(dataset["image"])
-    images = [i.convert("RGB") for i in images]
+    images = convert_if_not_rgb(images)
 
     start = time.time()
     line_predictions = batch_text_detection(images, det_model, det_processor)

benchmark/ordering.py

@@ -3,6 +3,7 @@
 import copy
 import json
 
+from surya.input.processing import convert_if_not_rgb
 from surya.model.ordering.model import load_model
 from surya.model.ordering.processor import load_processor
 from surya.ordering import batch_ordering
@@ -29,7 +30,7 @@ def main():
         split = f"train[:{args.max}]"
     dataset = datasets.load_dataset(settings.ORDER_BENCH_DATASET_NAME, split=split)
     images = list(dataset["image"])
-    images = [i.convert("RGB") for i in images]
+    images = convert_if_not_rgb(images)
     bboxes = list(dataset["bboxes"])
 
     start = time.time()

benchmark/recognition.py

@@ -4,6 +4,7 @@
 import torch
 
 from benchmark.scoring import overlap_score
+from surya.input.processing import convert_if_not_rgb
 from surya.model.recognition.model import load_model as load_recognition_model
 from surya.model.recognition.processor import load_processor as load_recognition_processor
 from surya.ocr import run_recognition
@@ -48,7 +49,7 @@ def main():
         dataset = dataset.filter(lambda x: x["language"] in langs)
 
     images = list(dataset["image"])
-    images = [i.convert("RGB") for i in images]
+    images = convert_if_not_rgb(images)
     bboxes = list(dataset["bboxes"])
     line_text = list(dataset["text"])
     languages = list(dataset["language"])

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "surya-ocr"
-version = "0.4.10"
+version = "0.4.11"
 description = "OCR, layout, reading order, and line detection in 90+ languages"
 authors = ["Vik Paruchuri <[email protected]>"]
 readme = "README.md"

surya/detection.py

@@ -7,7 +7,7 @@
 from surya.model.detection.segformer import SegformerForRegressionMask
 from surya.postprocessing.heatmap import get_and_clean_boxes
 from surya.postprocessing.affinity import get_vertical_lines
-from surya.input.processing import prepare_image_detection, split_image, get_total_splits
+from surya.input.processing import prepare_image_detection, split_image, get_total_splits, convert_if_not_rgb
 from surya.schema import TextDetectionResult
 from surya.settings import settings
 from tqdm import tqdm
@@ -51,7 +51,7 @@ def batch_detection(images: List, model: SegformerForRegressionMask, processor,
     all_preds = []
     for batch_idx in tqdm(range(len(batches)), desc="Detecting bboxes"):
         batch_image_idxs = batches[batch_idx]
-        batch_images = [images[j].convert("RGB") for j in batch_image_idxs]
+        batch_images = convert_if_not_rgb([images[j] for j in batch_image_idxs])
 
         split_index = []
         split_heights = []

surya/input/processing.py

@@ -1,5 +1,3 @@
-import os
-import random
 from typing import List
 
 import cv2
@@ -11,6 +9,15 @@
 from surya.settings import settings
 
 
+def convert_if_not_rgb(images: List[Image.Image]) -> List[Image.Image]:
+    new_images = []
+    for image in images:
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        new_images.append(image)
+    return new_images
+
+
 def get_total_splits(image_size, processor):
     img_height = list(image_size)[1]
     max_height = settings.DETECTOR_IMAGE_CHUNK_HEIGHT
@@ -48,6 +55,8 @@ def split_image(img, processor):
 def prepare_image_detection(img, processor):
     new_size = (processor.size["width"], processor.size["height"])
 
+    # This double resize actually necessary for downstream accuracy
+    img.thumbnail(new_size, Image.Resampling.LANCZOS)
     img = img.resize(new_size, Image.Resampling.LANCZOS) # Stretch smaller dimension to fit new size
 
     img = np.asarray(img, dtype=np.uint8)

surya/ocr.py

@@ -2,13 +2,14 @@
 from PIL import Image
 
 from surya.detection import batch_text_detection
-from surya.input.processing import slice_polys_from_image, slice_bboxes_from_image
+from surya.input.processing import slice_polys_from_image, slice_bboxes_from_image, convert_if_not_rgb
 from surya.postprocessing.text import sort_text_lines
 from surya.recognition import batch_recognition
 from surya.schema import TextLine, OCRResult
 
 
 def run_recognition(images: List[Image.Image], langs: List[List[str]], rec_model, rec_processor, bboxes: List[List[List[int]]] = None, polygons: List[List[List[List[int]]]] = None, batch_size=None) -> List[OCRResult]:
+    images = convert_if_not_rgb(images)
     # Polygons need to be in corner format - [[x1, y1], [x2, y2], [x3, y3], [x4, y4]], bboxes in [x1, y1, x2, y2] format
     assert bboxes is not None or polygons is not None
     assert len(images) == len(langs), "You need to pass in one list of languages for each image"
@@ -57,6 +58,7 @@ def run_recognition(images: List[Image.Image], langs: List[List[str]], rec_model
 
 
 def run_ocr(images: List[Image.Image], langs: List[List[str]], det_model, det_processor, rec_model, rec_processor, batch_size=None) -> List[OCRResult]:
+    images = convert_if_not_rgb(images)
     det_predictions = batch_text_detection(images, det_model, det_processor)
 
     all_slices = []

surya/ordering.py

@@ -3,6 +3,7 @@
 import torch
 from PIL import Image
 
+from surya.input.processing import convert_if_not_rgb
 from surya.model.ordering.encoderdecoder import OrderVisionEncoderDecoderModel
 from surya.schema import OrderBox, OrderResult
 from surya.settings import settings
@@ -37,7 +38,7 @@ def batch_ordering(images: List, bboxes: List[List[List[float]]], model: OrderVi
     if batch_size is None:
         batch_size = get_batch_size()
 
-    images = [image.convert("RGB") for image in images]
+    images = convert_if_not_rgb(images)
 
     output_order = []
     for i in tqdm(range(0, len(images), batch_size), desc="Finding reading order"):

surya/recognition.py

@@ -2,6 +2,7 @@
 import torch
 from PIL import Image
 
+from surya.input.processing import convert_if_not_rgb
 from surya.postprocessing.math.latex import fix_math, contains_math
 from surya.postprocessing.text import truncate_repetitions
 from surya.settings import settings
@@ -24,9 +25,11 @@ def get_batch_size():
 def batch_recognition(images: List, languages: List[List[str]], model, processor, batch_size=None):
     assert all([isinstance(image, Image.Image) for image in images])
     assert len(images) == len(languages)
-    assert [len(l) <= settings.RECOGNITION_MAX_LANGS for l in languages], f"OCR only supports up to {settings.RECOGNITION_MAX_LANGS} languages per image"
 
-    images = [image.convert("RGB") for image in images]
+    for l in languages:
+        assert len(l) <= settings.RECOGNITION_MAX_LANGS, f"OCR only supports up to {settings.RECOGNITION_MAX_LANGS} languages per image, you passed {l}."
+
+    images = convert_if_not_rgb(images)
     if batch_size is None:
         batch_size = get_batch_size()