nkozobrod/surya-ocr
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Suya OCR API — vLLM-backed, OpenAI-compatible OCR service

Browse Source 
FastAPI service wrapping the Surya-OCR-2 model (datalab-to) served through vLLM:
legacy /v1/api/ai/* endpoints, an OpenAI-compatible /v1/chat/completions endpoint,
a coalescing request batcher, a local OCR CLI, Docker packaging, multilingual
example outputs, and quantization/concurrency benchmarks.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Fu Dai
2026-06-17 10:20:02 +04:00
co-authored by Claude Opus 4.8
commit 1a585693be
147 changed files with 13827 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
surya/detection/__init__.py
+147
View File
@@ -0,0 +1,147 @@
from concurrent.futures import ThreadPoolExecutor
from typing import List, Generator, Tuple
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from tqdm import tqdm
from surya.common.predictor import BasePredictor
from surya.detection.loader import DetectionModelLoader
from surya.detection.parallel import FakeExecutor
from surya.detection.util import get_total_splits, split_image
from surya.detection.schema import TextDetectionResult
from surya.settings import settings
from surya.detection.heatmap import parallel_get_boxes
class DetectionPredictor(BasePredictor):
model_loader_cls = DetectionModelLoader
batch_size = settings.DETECTOR_BATCH_SIZE
default_batch_sizes = {"cpu": 8, "mps": 8, "cuda": 36}
def __call__(
self, images: List[Image.Image], batch_size=None, include_maps=False
) -> List[TextDetectionResult]:
detection_generator = self.batch_detection(images, batch_size=batch_size)
postprocessing_futures = []
max_workers = min(settings.DETECTOR_POSTPROCESSING_CPU_WORKERS, len(images))
parallelize = (
not settings.IN_STREAMLIT
and len(images) >= settings.DETECTOR_MIN_PARALLEL_THRESH
)
executor = ThreadPoolExecutor if parallelize else FakeExecutor
with executor(max_workers=max_workers) as e:
for preds, orig_sizes in detection_generator:
for pred, orig_size in zip(preds, orig_sizes):
postprocessing_futures.append(
e.submit(parallel_get_boxes, pred, orig_size, include_maps)
)
return [future.result() for future in postprocessing_futures]
def prepare_image(self, img):
new_size = (self.processor.size["width"], self.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)
img = self.processor(img)["pixel_values"][0]
img = torch.from_numpy(img)
return img
def batch_detection(
self, images: List, batch_size=None
) -> Generator[Tuple[List[List[np.ndarray]], List[Tuple[int, int]]], None, None]:
assert all([isinstance(image, Image.Image) for image in images])
if batch_size is None:
batch_size = self.get_batch_size()
heatmap_count = self.model.config.num_labels
orig_sizes = [image.size for image in images]
splits_per_image = [
get_total_splits(size, self.processor.size["height"]) for size in orig_sizes
]
batches = []
current_batch_size = 0
current_batch = []
for i in range(len(images)):
if current_batch_size + splits_per_image[i] > batch_size:
if len(current_batch) > 0:
batches.append(current_batch)
current_batch = []
current_batch_size = 0
current_batch.append(i)
current_batch_size += splits_per_image[i]
if len(current_batch) > 0:
batches.append(current_batch)
for batch_idx in tqdm(
range(len(batches)), desc="Detecting bboxes", disable=self.disable_tqdm
):
batch_image_idxs = batches[batch_idx]
batch_images = [images[j].convert("RGB") for j in batch_image_idxs]
split_index = []
split_heights = []
image_splits = []
for image_idx, image in enumerate(batch_images):
image_parts, split_height = split_image(
image, self.processor.size["height"]
)
image_splits.extend(image_parts)
split_index.extend([image_idx] * len(image_parts))
split_heights.extend(split_height)
image_splits = [self.prepare_image(image) for image in image_splits]
# Batch images in dim 0
batch = torch.stack(image_splits, dim=0).to(self.model.dtype)
with settings.INFERENCE_MODE():
pred = self.model(pixel_values=batch.to(self.model.device))
logits = pred.logits
correct_shape = [
self.processor.size["height"],
self.processor.size["width"],
]
current_shape = list(logits.shape[2:])
if current_shape != correct_shape:
logits = F.interpolate(
logits, size=correct_shape, mode="bilinear", align_corners=False
)
logits = logits.to(torch.float32).cpu().numpy()
preds = []
for i, (idx, height) in enumerate(zip(split_index, split_heights)):
# If our current prediction length is below the image idx, that means we have a new image
# Otherwise, we need to add to the current image
if len(preds) <= idx:
preds.append([logits[i][k] for k in range(heatmap_count)])
else:
heatmaps = preds[idx]
pred_heatmaps = [logits[i][k] for k in range(heatmap_count)]
if height < self.processor.size["height"]:
# Cut off padding to get original height
pred_heatmaps = [
pred_heatmap[:height, :] for pred_heatmap in pred_heatmaps
]
for k in range(heatmap_count):
heatmaps[k] = np.vstack([heatmaps[k], pred_heatmaps[k]])
preds[idx] = heatmaps
yield preds, [orig_sizes[j] for j in batch_image_idxs]
torch.cuda.empty_cache()
surya/detection/heatmap.py
+165
View File
@@ -0,0 +1,165 @@
from typing import List
import cv2
import numpy as np
from PIL import Image
from surya.common.util import clean_boxes
from surya.detection import TextDetectionResult
from surya.common.polygon import PolygonBox
from surya.settings import settings
def get_dynamic_thresholds(linemap, text_threshold, low_text, typical_top10_avg=0.7):
# Find average intensity of top 10% pixels
flat_map = linemap.ravel()
top_10_count = int(len(flat_map) * 0.9)
avg_intensity = np.mean(np.partition(flat_map, top_10_count)[top_10_count:])
scaling_factor = np.clip(avg_intensity / typical_top10_avg, 0, 1) ** (1 / 2)
low_text = np.clip(low_text * scaling_factor, 0.1, 0.6)
text_threshold = np.clip(text_threshold * scaling_factor, 0.15, 0.8)
return text_threshold, low_text
def detect_boxes(linemap, text_threshold, low_text):
# From CRAFT - https://github.com/clovaai/CRAFT-pytorch
# Modified to return boxes and for speed, accuracy
img_h, img_w = linemap.shape
text_threshold, low_text = get_dynamic_thresholds(linemap, text_threshold, low_text)
text_score_comb = (linemap > low_text).astype(np.uint8)
label_count, labels, stats, centroids = cv2.connectedComponentsWithStats(
text_score_comb, connectivity=4
)
det = []
confidences = []
max_confidence = 0
for k in range(1, label_count):
# size filtering
size = stats[k, cv2.CC_STAT_AREA]
if size < 10:
continue
# make segmentation map
x, y, w, h = stats[
k,
[cv2.CC_STAT_LEFT, cv2.CC_STAT_TOP, cv2.CC_STAT_WIDTH, cv2.CC_STAT_HEIGHT],
]
try:
niter = int(np.sqrt(min(w, h)))
except ValueError:
niter = 0
buffer = 1
sx, sy = max(0, x - niter - buffer), max(0, y - niter - buffer)
ex, ey = min(img_w, x + w + niter + buffer), min(img_h, y + h + niter + buffer)
mask = labels[sy:ey, sx:ex] == k
selected_linemap = linemap[sy:ey, sx:ex][mask]
if selected_linemap.size == 0:
continue
line_max = np.max(selected_linemap)
# thresholding
if line_max < text_threshold:
continue
segmap = mask.astype(np.uint8)
ksize = buffer + niter
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (ksize, ksize))
selected_segmap = cv2.dilate(segmap, kernel)
# make box
y_inds, x_inds = np.nonzero(selected_segmap)
x_inds += sx
y_inds += sy
np_contours = np.column_stack((x_inds, y_inds))
rectangle = cv2.minAreaRect(np_contours)
box = cv2.boxPoints(rectangle)
# align diamond-shape
w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
box_ratio = max(w, h) / (min(w, h) + 1e-5)
if abs(1 - box_ratio) <= 0.1:
left, right = np_contours[:, 0].min(), np_contours[:, 0].max()
top, bottom = np_contours[:, 1].min(), np_contours[:, 1].max()
box = np.array(
[[left, top], [right, top], [right, bottom], [left, bottom]],
dtype=np.float32,
)
# make clock-wise order
startidx = box.sum(axis=1).argmin()
box = np.roll(box, 4 - startidx, 0)
max_confidence = max(max_confidence, line_max)
confidences.append(line_max)
det.append(box)
if max_confidence > 0:
confidences = [c / max_confidence for c in confidences]
return det, confidences
def get_detected_boxes(textmap, text_threshold=None, low_text=None) -> List[PolygonBox]:
if text_threshold is None:
text_threshold = settings.DETECTOR_TEXT_THRESHOLD
if low_text is None:
low_text = settings.DETECTOR_BLANK_THRESHOLD
if textmap.dtype != np.float32:
textmap = textmap.astype(np.float32)
boxes, confidences = detect_boxes(textmap, text_threshold, low_text)
# From point form to box form
return [
PolygonBox(polygon=box, confidence=confidence)
for box, confidence in zip(boxes, confidences)
]
def get_and_clean_boxes(
textmap, processor_size, image_size, text_threshold=None, low_text=None
) -> List[PolygonBox]:
bboxes = get_detected_boxes(textmap, text_threshold, low_text)
for bbox in bboxes:
bbox.rescale(processor_size, image_size)
bbox.fit_to_bounds([0, 0, image_size[0], image_size[1]])
bboxes = clean_boxes(bboxes)
return bboxes
def parallel_get_boxes(preds, orig_sizes, include_maps=False):
heatmap, affinity_map = preds
heat_img, aff_img = None, None
if include_maps:
heat_img = Image.fromarray((heatmap * 255).astype(np.uint8))
aff_img = Image.fromarray((affinity_map * 255).astype(np.uint8))
heatmap_size = list(reversed(heatmap.shape))
bboxes = get_and_clean_boxes(heatmap, heatmap_size, orig_sizes)
for box in bboxes:
# Skip for vertical boxes
if box.height < 3 * box.width:
box.expand(x_margin=0, y_margin=settings.DETECTOR_BOX_Y_EXPAND_MARGIN)
box.fit_to_bounds(
[0, 0, orig_sizes[0], orig_sizes[1]]
) # Fix any bad expands
result = TextDetectionResult(
bboxes=bboxes,
heatmap=heat_img,
affinity_map=aff_img,
image_bbox=[0, 0, orig_sizes[0], orig_sizes[1]],
)
return result
surya/detection/loader.py
+53
View File
@@ -0,0 +1,53 @@
from typing import Optional
import torch
from surya.common.load import ModelLoader
from surya.detection.processor import SegformerImageProcessor
from surya.detection.model.config import EfficientViTConfig
from surya.detection.model.encoderdecoder import EfficientViTForSemanticSegmentation
from surya.logging import get_logger
from surya.settings import settings
logger = get_logger()
class DetectionModelLoader(ModelLoader):
def __init__(self, checkpoint: Optional[str] = None):
super().__init__(checkpoint)
if self.checkpoint is None:
self.checkpoint = settings.DETECTOR_MODEL_CHECKPOINT
def model(
self,
device: Optional[torch.device | str] = None,
dtype: Optional[torch.dtype | str] = None,
attention_implementation: Optional[str] = None,
) -> EfficientViTForSemanticSegmentation:
if device is None:
device = settings.TORCH_DEVICE_MODEL
if dtype is None:
dtype = settings.MODEL_DTYPE
config = EfficientViTConfig.from_pretrained(self.checkpoint)
model = EfficientViTForSemanticSegmentation.from_pretrained(
self.checkpoint,
dtype=dtype,
config=config,
)
model = model.to(device)
model = model.eval()
logger.debug(
f"Loaded detection model {self.checkpoint} from {EfficientViTForSemanticSegmentation.get_local_path(self.checkpoint)} onto device {device} with dtype {dtype}"
)
return model
def processor(
self,
device: Optional[torch.device | str] = None,
dtype: Optional[torch.dtype | str] = None,
) -> SegformerImageProcessor:
return SegformerImageProcessor.from_pretrained(self.checkpoint)
surya/detection/model/__init__.py
View File
surya/detection/model/config.py
+53
View File
@@ -0,0 +1,53 @@
from transformers import PretrainedConfig
from surya.common.s3 import S3DownloaderMixin
class EfficientViTConfig(S3DownloaderMixin, PretrainedConfig):
r"""
```"""
model_type = "efficientvit"
def __init__(
self,
num_classes=2,
num_channels=3,
widths=(32, 64, 128, 256, 512),
head_dim=32,
num_stages=4,
depths=(1, 1, 1, 6, 6),
strides=(2, 2, 2, 2, 2),
hidden_sizes=(32, 64, 160, 256),
patch_size=(7, 7),
hidden_dropout_prob=0.0,
attention_probs_dropout_prob=0.0,
classifier_dropout_prob=0.0,
layer_norm_eps=1e-6,
decoder_layer_hidden_size=128,
decoder_hidden_size=512,
semantic_loss_ignore_index=255,
initializer_range=0.02,
**kwargs,
):
super().__init__(**kwargs)
self.num_classes = num_classes
self.widths = widths
self.head_dim = head_dim
self.num_channels = num_channels
self.num_stages = num_stages
self.depths = depths
self.strides = strides
self.hidden_sizes = hidden_sizes
self.patch_size = patch_size
self.hidden_dropout_prob = hidden_dropout_prob
self.attention_probs_dropout_prob = attention_probs_dropout_prob
self.classifier_dropout_prob = classifier_dropout_prob
self.layer_norm_eps = layer_norm_eps
self.decoder_hidden_size = decoder_hidden_size
self.decoder_layer_hidden_size = decoder_layer_hidden_size
self.semantic_loss_ignore_index = semantic_loss_ignore_index
self.initializer_range = initializer_range
surya/detection/model/encoderdecoder.py
+839
View File
@@ -0,0 +1,839 @@
"""
This is an implementation of efficientvit, with some modifications (decode head, etc).
Original paper at https://arxiv.org/abs/2205.14756
Code adapted from timm, https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/efficientvit_mit.py
Original code (that timm adapted from) at https://github.com/mit-han-lab/efficientvit
License: Apache 2
"""
from __future__ import annotations
from typing import Optional, Union, Tuple, List, Any
from functools import partial
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.modeling_outputs import SemanticSegmenterOutput
from surya.common.pretrained import SuryaPreTrainedModel
from surya.common.s3 import S3DownloaderMixin
from surya.detection.model.config import EfficientViTConfig
def val2list(x: Union[List, Tuple, Any], repeat_time=1):
if isinstance(x, (list, tuple)):
return list(x)
return [x for _ in range(repeat_time)]
def val2tuple(x: Union[List, Tuple, Any], min_len: int = 1, idx_repeat: int = -1):
# repeat elements if necessary
x = val2list(x)
if len(x) > 0:
x[idx_repeat:idx_repeat] = [x[idx_repeat] for _ in range(min_len - len(x))]
return tuple(x)
def get_same_padding(
kernel_size: Union[int, Tuple[int, ...]],
) -> Union[int, Tuple[int, ...]]:
if isinstance(kernel_size, tuple):
return tuple([get_same_padding(ks) for ks in kernel_size])
else:
assert kernel_size % 2 > 0, "kernel size should be odd number"
return kernel_size // 2
def get_padding(kernel_size: int, stride: int = 1, dilation: int = 1) -> int:
padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
return padding
class ConvNormAct(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size=3,
stride=1,
dilation=1,
groups=1,
bias=False,
dropout=0.0,
norm_layer=nn.BatchNorm2d,
act_layer=nn.ReLU,
):
super(ConvNormAct, self).__init__()
self.dropout = nn.Dropout(dropout, inplace=False)
padding = get_padding(kernel_size, stride, dilation)
self.conv = nn.Conv2d(
in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
groups=groups,
bias=bias,
padding=padding,
)
self.norm = (
norm_layer(num_features=out_channels) if norm_layer else nn.Identity()
)
self.act = act_layer(inplace=True) if act_layer is not None else nn.Identity()
def forward(self, x):
x = self.conv(x)
x = self.norm(x)
x = self.act(x)
return x
class DSConv(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size=3,
stride=1,
use_bias=False,
norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d),
act_layer=(nn.ReLU6, None),
):
super(DSConv, self).__init__()
use_bias = val2tuple(use_bias, 2)
norm_layer = val2tuple(norm_layer, 2)
act_layer = val2tuple(act_layer, 2)
self.depth_conv = ConvNormAct(
in_channels,
in_channels,
kernel_size,
stride,
groups=in_channels,
norm_layer=norm_layer[0],
act_layer=act_layer[0],
bias=use_bias[0],
)
self.point_conv = ConvNormAct(
in_channels,
out_channels,
1,
norm_layer=norm_layer[1],
act_layer=act_layer[1],
bias=use_bias[1],
)
def forward(self, x):
x = self.depth_conv(x)
x = self.point_conv(x)
return x
class ConvBlock(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size=3,
stride=1,
mid_channels=None,
expand_ratio=1,
use_bias=False,
norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d),
act_layer=(nn.ReLU6, None),
):
super(ConvBlock, self).__init__()
use_bias = val2tuple(use_bias, 2)
norm_layer = val2tuple(norm_layer, 2)
act_layer = val2tuple(act_layer, 2)
mid_channels = mid_channels or round(in_channels * expand_ratio)
self.conv1 = ConvNormAct(
in_channels,
mid_channels,
kernel_size,
stride,
norm_layer=norm_layer[0],
act_layer=act_layer[0],
bias=use_bias[0],
)
self.conv2 = ConvNormAct(
mid_channels,
out_channels,
kernel_size,
1,
norm_layer=norm_layer[1],
act_layer=act_layer[1],
bias=use_bias[1],
)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
return x
class MBConv(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size=3,
stride=1,
mid_channels=None,
expand_ratio=6,
use_bias=False,
norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d, nn.BatchNorm2d),
act_layer=(nn.ReLU6, nn.ReLU6, None),
):
super(MBConv, self).__init__()
use_bias = val2tuple(use_bias, 3)
norm_layer = val2tuple(norm_layer, 3)
act_layer = val2tuple(act_layer, 3)
mid_channels = mid_channels or round(in_channels * expand_ratio)
self.inverted_conv = ConvNormAct(
in_channels,
mid_channels,
1,
stride=1,
norm_layer=norm_layer[0],
act_layer=act_layer[0],
bias=use_bias[0],
)
self.depth_conv = ConvNormAct(
mid_channels,
mid_channels,
kernel_size,
stride=stride,
groups=mid_channels,
norm_layer=norm_layer[1],
act_layer=act_layer[1],
bias=use_bias[1],
)
self.point_conv = ConvNormAct(
mid_channels,
out_channels,
1,
norm_layer=norm_layer[2],
act_layer=act_layer[2],
bias=use_bias[2],
)
def forward(self, x):
x = self.inverted_conv(x)
x = self.depth_conv(x)
x = self.point_conv(x)
return x
class FusedMBConv(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size=3,
stride=1,
mid_channels=None,
expand_ratio=6,
groups=1,
use_bias=False,
norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d),
act_layer=(nn.ReLU6, None),
):
super(FusedMBConv, self).__init__()
use_bias = val2tuple(use_bias, 2)
norm_layer = val2tuple(norm_layer, 2)
act_layer = val2tuple(act_layer, 2)
mid_channels = mid_channels or round(in_channels * expand_ratio)
self.spatial_conv = ConvNormAct(
in_channels,
mid_channels,
kernel_size,
stride=stride,
groups=groups,
norm_layer=norm_layer[0],
act_layer=act_layer[0],
bias=use_bias[0],
)
self.point_conv = ConvNormAct(
mid_channels,
out_channels,
1,
norm_layer=norm_layer[1],
act_layer=act_layer[1],
bias=use_bias[1],
)
def forward(self, x):
x = self.spatial_conv(x)
x = self.point_conv(x)
return x
class LiteMLA(nn.Module):
"""Lightweight multi-scale linear attention"""
def __init__(
self,
in_channels: int,
out_channels: int,
heads: Union[int, None] = None,
heads_ratio: float = 1.0,
dim=8,
use_bias=False,
norm_layer=(None, nn.BatchNorm2d),
act_layer=(None, None),
kernel_func=nn.ReLU,
scales=(5,),
eps=1e-5,
):
super(LiteMLA, self).__init__()
self.eps = eps
heads = heads or int(in_channels // dim * heads_ratio)
total_dim = heads * dim
use_bias = val2tuple(use_bias, 2)
norm_layer = val2tuple(norm_layer, 2)
act_layer = val2tuple(act_layer, 2)
self.dim = dim
self.qkv = ConvNormAct(
in_channels,
3 * total_dim,
1,
bias=use_bias[0],
norm_layer=norm_layer[0],
act_layer=act_layer[0],
)
self.aggreg = nn.ModuleList(
[
nn.Sequential(
nn.Conv2d(
3 * total_dim,
3 * total_dim,
scale,
padding=get_same_padding(scale),
groups=3 * total_dim,
bias=use_bias[0],
),
nn.Conv2d(
3 * total_dim,
3 * total_dim,
1,
groups=3 * heads,
bias=use_bias[0],
),
)
for scale in scales
]
)
self.kernel_func = kernel_func(inplace=False)
self.proj = ConvNormAct(
total_dim * (1 + len(scales)),
out_channels,
1,
bias=use_bias[1],
norm_layer=norm_layer[1],
act_layer=act_layer[1],
)
def _attn(self, q, k, v):
dtype = v.dtype
q, k, v = q.float(), k.float(), v.float()
kv = k.transpose(-1, -2) @ v
out = q @ kv
out = out[..., :-1] / (out[..., -1:] + self.eps)
return out.to(dtype)
def forward(self, x):
# Shape is B, C, H, W
B, _, H, W = x.shape
# generate multi-scale q, k, v
qkv = self.qkv(x)
multi_scale_qkv = [qkv]
for op in self.aggreg:
multi_scale_qkv.append(op(qkv))
multi_scale_qkv = torch.cat(multi_scale_qkv, dim=1)
multi_scale_qkv = multi_scale_qkv.reshape(B, -1, 3 * self.dim, H * W).transpose(
-1, -2
)
# Shape for each is B, C, HW, head_dim
q, k, v = multi_scale_qkv.chunk(3, dim=-1)
# lightweight global attention
q = self.kernel_func(q)
k = self.kernel_func(k)
v = F.pad(v, (0, 1), mode="constant", value=1.0)
out = self._attn(q, k, v)
# final projection
out = out.transpose(-1, -2).reshape(B, -1, H, W)
out = self.proj(out)
return out
class EfficientVitBlock(nn.Module):
def __init__(
self,
in_channels,
heads_ratio=1.0,
head_dim=32,
expand_ratio=4,
norm_layer=nn.BatchNorm2d,
act_layer=nn.Hardswish,
):
super(EfficientVitBlock, self).__init__()
self.context_module = ResidualBlock(
LiteMLA(
in_channels=in_channels,
out_channels=in_channels,
heads_ratio=heads_ratio,
dim=head_dim,
norm_layer=(None, norm_layer),
),
nn.Identity(),
)
self.local_module = ResidualBlock(
MBConv(
in_channels=in_channels,
out_channels=in_channels,
expand_ratio=expand_ratio,
use_bias=(True, True, False),
norm_layer=(None, None, norm_layer),
act_layer=(act_layer, act_layer, None),
),
nn.Identity(),
)
def forward(self, x):
x = self.context_module(x)
x = self.local_module(x)
return x
class ResidualBlock(nn.Module):
def __init__(
self,
main: Optional[nn.Module],
shortcut: Optional[nn.Module] = None,
pre_norm: Optional[nn.Module] = None,
):
super(ResidualBlock, self).__init__()
self.pre_norm = pre_norm if pre_norm is not None else nn.Identity()
self.main = main
self.shortcut = shortcut
def forward(self, x):
res = self.main(self.pre_norm(x))
if self.shortcut is not None:
res = res + self.shortcut(x)
return res
def build_local_block(
in_channels: int,
out_channels: int,
stride: int,
kernel_size: int,
expand_ratio: float,
norm_layer: str,
act_layer: str,
fewer_norm: bool = False,
block_type: str = "default",
):
assert block_type in ["default", "large", "fused"]
if expand_ratio == 1:
if block_type == "default":
block = DSConv(
in_channels=in_channels,
out_channels=out_channels,
stride=stride,
kernel_size=kernel_size,
use_bias=(True, False) if fewer_norm else False,
norm_layer=(None, norm_layer) if fewer_norm else norm_layer,
act_layer=(act_layer, None),
)
else:
block = ConvBlock(
in_channels=in_channels,
out_channels=out_channels,
stride=stride,
kernel_size=kernel_size,
use_bias=(True, False) if fewer_norm else False,
norm_layer=(None, norm_layer) if fewer_norm else norm_layer,
act_layer=(act_layer, None),
)
else:
if block_type == "default":
block = MBConv(
in_channels=in_channels,
out_channels=out_channels,
stride=stride,
kernel_size=kernel_size,
expand_ratio=expand_ratio,
use_bias=(True, True, False) if fewer_norm else False,
norm_layer=(None, None, norm_layer) if fewer_norm else norm_layer,
act_layer=(act_layer, act_layer, None),
)
else:
block = FusedMBConv(
in_channels=in_channels,
out_channels=out_channels,
stride=stride,
kernel_size=kernel_size,
expand_ratio=expand_ratio,
use_bias=(True, False) if fewer_norm else False,
norm_layer=(None, norm_layer) if fewer_norm else norm_layer,
act_layer=(act_layer, None),
)
return block
class Stem(nn.Sequential):
def __init__(
self,
in_chs,
out_chs,
depth,
stride,
norm_layer,
act_layer,
block_type="default",
):
super().__init__()
self.stride = stride
self.add_module(
"in_conv",
ConvNormAct(
in_chs,
out_chs,
kernel_size=stride + 1,
stride=stride,
norm_layer=norm_layer,
act_layer=act_layer,
),
)
stem_block = 0
for _ in range(depth):
self.add_module(
f"res{stem_block}",
ResidualBlock(
build_local_block(
in_channels=out_chs,
out_channels=out_chs,
stride=1,
kernel_size=3,
expand_ratio=1,
norm_layer=norm_layer,
act_layer=act_layer,
block_type=block_type,
),
nn.Identity(),
),
)
stem_block += 1
class EfficientVitLargeStage(nn.Module):
def __init__(
self,
in_chs,
out_chs,
depth,
stride,
norm_layer,
act_layer,
head_dim,
vit_stage=False,
fewer_norm=False,
):
super(EfficientVitLargeStage, self).__init__()
blocks = [
ResidualBlock(
build_local_block(
in_channels=in_chs,
out_channels=out_chs,
stride=stride,
kernel_size=stride + 1,
expand_ratio=24 if vit_stage else 16,
norm_layer=norm_layer,
act_layer=act_layer,
fewer_norm=vit_stage or fewer_norm,
block_type="default" if fewer_norm else "fused",
),
None,
)
]
in_chs = out_chs
if vit_stage:
# for stage 4
for _ in range(depth):
blocks.append(
EfficientVitBlock(
in_channels=in_chs,
head_dim=head_dim,
expand_ratio=6,
norm_layer=norm_layer,
act_layer=act_layer,
)
)
else:
# for stage 1, 2, 3
for i in range(depth):
blocks.append(
ResidualBlock(
build_local_block(
in_channels=in_chs,
out_channels=out_chs,
stride=1,
kernel_size=3,
expand_ratio=4,
norm_layer=norm_layer,
act_layer=act_layer,
fewer_norm=fewer_norm,
block_type="default" if fewer_norm else "fused",
),
nn.Identity(),
)
)
self.blocks = nn.Sequential(*blocks)
def forward(self, x):
return self.blocks(x)
class EfficientVitLarge(nn.Module):
def __init__(
self,
config: EfficientViTConfig,
norm_layer=nn.BatchNorm2d,
act_layer=nn.Hardswish,
):
super(EfficientVitLarge, self).__init__()
self.grad_checkpointing = False
self.num_classes = config.num_classes
self.norm_eps = config.layer_norm_eps
norm_layer = partial(norm_layer, eps=self.norm_eps)
# input stem
self.stem = Stem(
config.num_channels,
config.widths[0],
config.depths[0],
config.strides[0],
norm_layer,
act_layer,
block_type="large",
)
stride = config.strides[0]
# stages
self.feature_info = []
self.stages = nn.Sequential()
in_channels = config.widths[0]
for i, (w, d, s) in enumerate(
zip(config.widths[1:], config.depths[1:], config.strides[1:])
):
self.stages.append(
EfficientVitLargeStage(
in_channels,
w,
depth=d,
stride=s,
norm_layer=norm_layer,
act_layer=act_layer,
head_dim=config.head_dim,
vit_stage=i >= 3,
fewer_norm=i >= 2,
)
)
stride *= s
in_channels = w
self.feature_info += [
dict(num_chs=in_channels, reduction=stride, module=f"stages.{i}")
]
self.num_features = in_channels
@torch.jit.ignore
def set_grad_checkpointing(self, enable=True):
self.grad_checkpointing = enable
def forward(self, x):
x = self.stem(x)
encoder_hidden_states = []
for i, module in enumerate(self.stages):
x = module(x)
encoder_hidden_states.append(x)
return encoder_hidden_states
class EfficientViTPreTrainedModel(SuryaPreTrainedModel):
"""
An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
models.
"""
config_class = EfficientViTConfig
base_model_prefix = "efficientvit"
main_input_name = "pixel_values"
def _init_weights(self, module):
"""Initialize the weights"""
if isinstance(module, (nn.Linear, nn.Conv2d)):
# Slightly different from the TF version which uses truncated_normal for initialization
# cf https://github.com/pytorch/pytorch/pull/5617
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
if module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.Embedding):
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
if module.padding_idx is not None:
module.weight.data[module.padding_idx].zero_()
elif isinstance(module, nn.LayerNorm):
module.bias.data.zero_()
module.weight.data.fill_(1.0)
class DecodeMLP(nn.Module):
def __init__(self, input_dim, output_dim):
super().__init__()
self.proj = nn.Linear(input_dim, output_dim)
def forward(self, hidden_states: torch.Tensor):
# Input is B, C, H, W
hidden_states = hidden_states.flatten(2).transpose(1, 2)
# Output is B, HW, C
hidden_states = self.proj(hidden_states)
return hidden_states
class DecodeHead(EfficientViTPreTrainedModel):
def __init__(self, config: EfficientViTConfig):
super().__init__(config)
# linear layers which will unify the channel dimension of each of the encoder blocks to the same config.decoder_hidden_size
mlps = []
for width in config.widths[1:]:
mlp = DecodeMLP(
input_dim=width, output_dim=config.decoder_layer_hidden_size
)
mlps.append(mlp)
self.linear_c = nn.ModuleList(mlps)
# the following 3 layers implement the ConvModule of the original implementation
self.linear_fuse = nn.Conv2d(
in_channels=config.decoder_layer_hidden_size * config.num_stages,
out_channels=config.decoder_hidden_size,
kernel_size=1,
bias=False,
)
self.batch_norm = nn.BatchNorm2d(config.decoder_hidden_size)
self.activation = nn.ReLU()
self.dropout = nn.Dropout(config.classifier_dropout_prob)
self.classifier = nn.Conv2d(
config.decoder_hidden_size, config.num_labels, kernel_size=1
)
self.config = config
def forward(self, encoder_hidden_states: torch.FloatTensor) -> torch.Tensor:
batch_size = encoder_hidden_states[-1].shape[0]
all_hidden_states = ()
for encoder_hidden_state, mlp in zip(encoder_hidden_states, self.linear_c):
height, width = encoder_hidden_state.shape[2], encoder_hidden_state.shape[3]
encoder_hidden_state = mlp(encoder_hidden_state) # Output is B, HW, C
# Permute to B, C, HW
encoder_hidden_state = encoder_hidden_state.permute(0, 2, 1)
encoder_hidden_state = encoder_hidden_state.reshape(
batch_size, -1, height, width
)
# upsample
encoder_hidden_state = nn.functional.interpolate(
encoder_hidden_state,
size=encoder_hidden_states[0].size()[2:],
mode="bilinear",
align_corners=False,
)
all_hidden_states += (encoder_hidden_state,)
hidden_states = self.linear_fuse(torch.cat(all_hidden_states[::-1], dim=1))
hidden_states = self.batch_norm(hidden_states)
hidden_states = self.activation(hidden_states)
# logits are of shape (batch_size, num_labels, height/4, width/4)
logits = self.classifier(hidden_states)
return logits
class EfficientViTForSemanticSegmentation(
S3DownloaderMixin, EfficientViTPreTrainedModel
):
def __init__(self, config, **kwargs):
super().__init__(config)
self.vit = EfficientVitLarge(config)
self.decode_head = DecodeHead(config)
# Initialize weights and apply final processing
self.post_init()
def forward(
self, pixel_values: torch.FloatTensor
) -> Union[Tuple, SemanticSegmenterOutput]:
# Pixel values should be B,C,H,W
encoder_hidden_states = self.vit(
pixel_values,
)
logits = self.decode_head(encoder_hidden_states)
# Apply sigmoid to get 0-1 output
logits = torch.special.expit(logits)
return SemanticSegmenterOutput(
loss=None, logits=logits, hidden_states=encoder_hidden_states
)
class EfficientViTForSemanticLayoutSegmentation(EfficientViTPreTrainedModel):
def __init__(self, config, **kwargs):
super().__init__(config, **kwargs)
self.vit = EfficientVitLarge(config)
self.decode_head = DecodeHead(config)
# Initialize weights and apply final processing
self.post_init()
def forward(
self, pixel_values: torch.FloatTensor
) -> Union[Tuple, SemanticSegmenterOutput]:
# Pixel values should be B,C,H,W
encoder_hidden_states = self.vit(
pixel_values,
)
logits = self.decode_head(encoder_hidden_states)
# Apply sigmoid to get 0-1 output
logits = torch.special.expit(logits)
return SemanticSegmenterOutput(
loss=None, logits=logits, hidden_states=encoder_hidden_states
)
surya/detection/parallel.py
+19
View File
@@ -0,0 +1,19 @@
class FakeFuture:
def __init__(self, func, *args, **kwargs):
self._result = func(*args, **kwargs)
def result(self):
return self._result
class FakeExecutor:
def __init__(self, **kwargs):
pass
def __enter__(self):
return self
def __exit__(self, *excinfo):
pass
def submit(self, fn, *args, **kwargs):
return FakeFuture(fn, *args, **kwargs)
surya/detection/processor.py
+317
View File
@@ -0,0 +1,317 @@
# coding=utf-8
# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Modified image processor class for Segformer based on transformers"""
import warnings
from typing import Any, Dict, List, Optional, Union
import numpy as np
from transformers.image_processing_utils import (
BaseImageProcessor,
BatchFeature,
get_size_dict,
)
from transformers.image_transforms import to_channel_dimension_format
from transformers.image_utils import (
IMAGENET_DEFAULT_MEAN,
IMAGENET_DEFAULT_STD,
ChannelDimension,
ImageInput,
PILImageResampling,
infer_channel_dimension_format,
make_list_of_images,
)
from transformers.utils import TensorType
import PIL.Image
from surya.common.s3 import S3DownloaderMixin
class SegformerImageProcessor(S3DownloaderMixin, BaseImageProcessor):
r"""
Constructs a Segformer image processor.
Args:
do_resize (`bool`, *optional*, defaults to `True`):
Whether to resize the image's (height, width) dimensions to the specified `(size["height"],
size["width"])`. Can be overridden by the `do_resize` parameter in the `preprocess` method.
size (`Dict[str, int]` *optional*, defaults to `{"height": 512, "width": 512}`):
Size of the output image after resizing. Can be overridden by the `size` parameter in the `preprocess`
method.
resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):
Resampling filter to use if resizing the image. Can be overridden by the `resample` parameter in the
`preprocess` method.
do_rescale (`bool`, *optional*, defaults to `True`):
Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the `do_rescale`
parameter in the `preprocess` method.
rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
method.
do_normalize (`bool`, *optional*, defaults to `True`):
Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
method.
image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
Mean to use if normalizing the image. This is a float or list of floats the length of the number of
channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
do_reduce_labels (`bool`, *optional*, defaults to `False`):
Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0 is
used for background, and background itself is not included in all classes of a dataset (e.g. ADE20k). The
background label will be replaced by 255. Can be overridden by the `do_reduce_labels` parameter in the
`preprocess` method.
"""
model_input_names = ["pixel_values"]
def __init__(
self,
do_resize: bool = True,
size: Dict[str, int] = None,
resample: PILImageResampling = PILImageResampling.BILINEAR,
do_rescale: bool = True,
rescale_factor: Union[int, float] = 1 / 255,
do_normalize: bool = True,
image_mean: Optional[Union[float, List[float]]] = None,
image_std: Optional[Union[float, List[float]]] = None,
do_reduce_labels: bool = False,
**kwargs,
) -> None:
if "reduce_labels" in kwargs:
warnings.warn(
"The `reduce_labels` parameter is deprecated and will be removed in a future version. Please use "
"`do_reduce_labels` instead.",
FutureWarning,
)
do_reduce_labels = kwargs.pop("reduce_labels")
super().__init__(**kwargs)
size = size if size is not None else {"height": 512, "width": 512}
size = get_size_dict(size)
self.do_resize = do_resize
self.size = size
self.resample = resample
self.do_rescale = do_rescale
self.rescale_factor = rescale_factor
self.do_normalize = do_normalize
self.image_mean = (
image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
)
self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
self.do_reduce_labels = do_reduce_labels
self._valid_processor_keys = [
"images",
"segmentation_maps",
"do_resize",
"size",
"resample",
"do_rescale",
"rescale_factor",
"do_normalize",
"image_mean",
"image_std",
"do_reduce_labels",
"return_tensors",
"data_format",
"input_data_format",
]
@classmethod
def from_dict(cls, image_processor_dict: Dict[str, Any], **kwargs):
"""
Overrides the `from_dict` method from the base class to make sure `do_reduce_labels` is updated if image
processor is created using from_dict and kwargs e.g. `SegformerImageProcessor.from_pretrained(checkpoint,
reduce_labels=True)`
"""
image_processor_dict = image_processor_dict.copy()
if "reduce_labels" in kwargs:
image_processor_dict["reduce_labels"] = kwargs.pop("reduce_labels")
return super().from_dict(image_processor_dict, **kwargs)
def _preprocess(
self,
image: ImageInput,
do_resize: bool,
do_rescale: bool,
do_normalize: bool,
size: Optional[Dict[str, int]] = None,
resample: PILImageResampling = None,
rescale_factor: Optional[float] = None,
image_mean: Optional[Union[float, List[float]]] = None,
image_std: Optional[Union[float, List[float]]] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
):
if do_rescale:
image = self.rescale(
image=image, scale=rescale_factor, input_data_format=input_data_format
)
if do_normalize:
image = self.normalize(
image=image,
mean=image_mean,
std=image_std,
input_data_format=input_data_format,
)
return image
def _preprocess_image(
self,
image: ImageInput,
do_resize: bool = None,
size: Dict[str, int] = None,
resample: PILImageResampling = None,
do_rescale: bool = None,
rescale_factor: float = None,
do_normalize: bool = None,
image_mean: Optional[Union[float, List[float]]] = None,
image_std: Optional[Union[float, List[float]]] = None,
data_format: Optional[Union[str, ChannelDimension]] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
) -> np.ndarray:
"""Preprocesses a single image."""
# All transformations expect numpy arrays.
if input_data_format is None:
input_data_format = infer_channel_dimension_format(image)
image = self._preprocess(
image=image,
do_resize=do_resize,
size=size,
resample=resample,
do_rescale=do_rescale,
rescale_factor=rescale_factor,
do_normalize=do_normalize,
image_mean=image_mean,
image_std=image_std,
input_data_format=input_data_format,
)
if data_format is not None:
image = to_channel_dimension_format(
image, data_format, input_channel_dim=input_data_format
)
return image
def __call__(self, images, segmentation_maps=None, **kwargs):
"""
Preprocesses a batch of images and optionally segmentation maps.
Overrides the `__call__` method of the `Preprocessor` class so that both images and segmentation maps can be
passed in as positional arguments.
"""
return super().__call__(images, segmentation_maps=segmentation_maps, **kwargs)
def preprocess(
self,
images: ImageInput,
segmentation_maps: Optional[ImageInput] = None,
do_resize: Optional[bool] = None,
size: Optional[Dict[str, int]] = None,
resample: PILImageResampling = None,
do_rescale: Optional[bool] = None,
rescale_factor: Optional[float] = None,
do_normalize: Optional[bool] = None,
image_mean: Optional[Union[float, List[float]]] = None,
image_std: Optional[Union[float, List[float]]] = None,
do_reduce_labels: Optional[bool] = None,
return_tensors: Optional[Union[str, TensorType]] = None,
data_format: ChannelDimension = ChannelDimension.FIRST,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
**kwargs,
) -> PIL.Image.Image:
"""
Preprocess an image or batch of images.
Args:
images (`ImageInput`):
Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
passing in images with pixel values between 0 and 1, set `do_rescale=False`.
segmentation_maps (`ImageInput`, *optional*):
Segmentation map to preprocess.
do_resize (`bool`, *optional*, defaults to `self.do_resize`):
Whether to resize the image.
size (`Dict[str, int]`, *optional*, defaults to `self.size`):
Size of the image after `resize` is applied.
resample (`int`, *optional*, defaults to `self.resample`):
Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`, Only
has an effect if `do_resize` is set to `True`.
do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
Whether to rescale the image values between [0 - 1].
rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
Rescale factor to rescale the image by if `do_rescale` is set to `True`.
do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
Whether to normalize the image.
image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
Image mean.
image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
Image standard deviation.
do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`):
Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0
is used for background, and background itself is not included in all classes of a dataset (e.g.
ADE20k). The background label will be replaced by 255.
return_tensors (`str` or `TensorType`, *optional*):
The type of tensors to return. Can be one of:
- Unset: Return a list of `np.ndarray`.
- `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
The channel dimension format for the output image. Can be one of:
- `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- `ChannelDimension.LAST`: image in (height, width, num_channels) format.
input_data_format (`ChannelDimension` or `str`, *optional*):
The channel dimension format for the input image. If unset, the channel dimension format is inferred
from the input image. Can be one of:
- `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
"""
do_resize = do_resize if do_resize is not None else self.do_resize
do_rescale = do_rescale if do_rescale is not None else self.do_rescale
do_normalize = do_normalize if do_normalize is not None else self.do_normalize
resample = resample if resample is not None else self.resample
size = size if size is not None else self.size
rescale_factor = (
rescale_factor if rescale_factor is not None else self.rescale_factor
)
image_mean = image_mean if image_mean is not None else self.image_mean
image_std = image_std if image_std is not None else self.image_std
images = make_list_of_images(images)
images = [
self._preprocess_image(
image=img,
do_resize=do_resize,
resample=resample,
size=size,
do_rescale=do_rescale,
rescale_factor=rescale_factor,
do_normalize=do_normalize,
image_mean=image_mean,
image_std=image_std,
data_format=data_format,
input_data_format=input_data_format,
)
for img in images
]
data = {"pixel_values": images}
return BatchFeature(data=data, tensor_type=return_tensors)
surya/detection/schema.py
+12
View File
@@ -0,0 +1,12 @@
from typing import List, Optional, Any
from pydantic import BaseModel
from surya.common.polygon import PolygonBox
class TextDetectionResult(BaseModel):
bboxes: List[PolygonBox]
heatmap: Optional[Any]
affinity_map: Optional[Any]
image_bbox: List[float]
surya/detection/util.py
+36
View File
@@ -0,0 +1,36 @@
import math
from PIL import ImageOps
from surya.settings import settings
def get_total_splits(image_size, height):
img_height = list(image_size)[1]
max_height = settings.DETECTOR_IMAGE_CHUNK_HEIGHT
if img_height > max_height:
num_splits = math.ceil(img_height / height)
return num_splits
return 1
def split_image(img, height):
# This will not modify/return the original image - it will either crop, or copy the image
img_height = list(img.size)[1]
max_height = settings.DETECTOR_IMAGE_CHUNK_HEIGHT
if img_height > max_height:
num_splits = math.ceil(img_height / height)
splits = []
split_heights = []
for i in range(num_splits):
top = i * height
bottom = (i + 1) * height
if bottom > img_height:
bottom = img_height
cropped = img.crop((0, top, img.size[0], bottom))
chunk_height = bottom - top
if chunk_height < height:
cropped = ImageOps.pad(cropped, (img.size[0], height), color=255, centering=(0, 0))
splits.append(cropped)
split_heights.append(chunk_height)
return splits, split_heights
return [img.copy()], [img_height]