Merge pull request #120 from VikParuchuri/dev

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)