File size: 10,496 Bytes
4f7bf82 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 | import torch
import torch.nn.functional as F
import cv2
from PIL import Image
from torch.utils.data import DataLoader
from .dataset import Makeset
def generate(model, image, bos_id, eos_id, device="cuda", temp=0.5, max_iter=64, penalty=1.15, top_k=5):
"""
Prediction function using greedy search / top-k sampling with repetition penalty.
Args:
model: SOCRATE model instance.
image (Tensor): Single image tensor [1, C, H, W].
bos_id (int): Begin-of-sequence token ID.
eos_id (int): End-of-sequence token ID.
device (str): Target device. Default: "cuda".
temp (float): Temperature for sampling. Lower = more greedy. Default: 0.5.
max_iter (int): Max number of tokens to generate. Default: 64.
penalty (float): Repetition penalty applied to already-seen tokens. Default: 1.15.
top_k (int): Number of top candidates to sample from at each step. Default: 5.
"""
model.eval()
current_text = [bos_id]
generated = []
already_seen = set()
with torch.inference_mode():
memory_image = model.encode(image)
for i in range(max_iter):
x = torch.tensor([current_text], dtype=torch.long).to(device)
output = model.decode(memory_image, x)
output = output[:, -1, :]
for token_id in already_seen:
if output[0, token_id] < 0:
output[0, token_id] *= penalty
else:
output[0, token_id] /= penalty
output = output / temp
topk_vals, topk_idx = torch.topk(output, top_k, dim=-1)
probs = F.softmax(topk_vals, dim=-1)
idx = torch.multinomial(probs, 1)
idx = topk_idx.gather(-1, idx).item()
if idx == eos_id:
break
generated.append(idx)
current_text.append(idx)
already_seen.add(idx)
return generated
def generate_fast(model, image, bos_id, eos_id, device="cuda", max_iter=32):
"""
Super-fast prediction using only argmax (no sampling).
Args:
model: SOCRATE model instance.
image (Tensor): Single image tensor [1, C, H, W].
bos_id (int): Begin-of-sequence token ID.
eos_id (int): End-of-sequence token ID.
device (str): Target device. Default: "cuda".
max_iter (int): Max number of tokens to generate. Default: 32.
"""
model.eval()
current_text = [bos_id]
generated = []
with torch.inference_mode():
memory_image = model.encode(image)
for _ in range(max_iter):
x = torch.tensor([current_text], dtype=torch.long, device=device)
output = model.decode(memory_image, x)
logits = output[:, -1, :]
idx = logits.argmax(dim=-1).item()
if idx == eos_id:
break
generated.append(idx)
current_text.append(idx)
return generated
def beam_search(model, image, bos_id, eos_id, device="cuda", beam_width=4, max_iter=64):
"""
Beam search decoding.
Args:
model: SOCRATE model instance.
image (Tensor): Single image tensor [1, C, H, W].
bos_id (int): Begin-of-sequence token ID.
eos_id (int): End-of-sequence token ID.
device (str): Target device. Default: "cuda".
beam_width (int): Number of beams. Default: 4.
max_iter (int): Max tokens per beam. Default: 64.
Note: Full beam search is coming soon. Currently uses generate_fast as a fallback.
"""
print("WARNING: Beam search is not fully implemented yet. Using generate_fast as a fallback.")
return generate_fast(model, image, bos_id, eos_id, device, max_iter=max_iter)
def extract_crops_from_image(image_path, doctr_model=None):
"""
Extracts words (crops) using doctr and sorts them
correctly from top-to-bottom and left-to-right.
"""
if doctr_model is None:
from doctr.models import detection_predictor
doctr_model = detection_predictor(arch="db_resnet50", pretrained=True)
from doctr.io import DocumentFile
doc = DocumentFile.from_images(image_path)
result = doctr_model(doc)
boxes = result[0]["words"]
image = cv2.imread(image_path)
if image is None:
raise ValueError(f"Could not load image from {image_path}")
H, W = image.shape[:2]
# Sort by lines a la SOCRATE
boxes_info = []
for b in boxes:
xmin, ymin, xmax, ymax, score = b
cy = (ymin + ymax) / 2.0
h = ymax - ymin
boxes_info.append({'box': b, 'cy': cy, 'h': h, 'x': xmin})
boxes_info.sort(key=lambda item: item['cy'])
lines = []
current_line = []
for b in boxes_info:
if not current_line:
current_line.append(b)
else:
tolerance = current_line[0]['h'] * 0.5
if abs(b['cy'] - current_line[0]['cy']) < tolerance:
current_line.append(b)
else:
lines.append(current_line)
current_line = [b]
if current_line:
lines.append(current_line)
sorted_boxes = []
for line in lines:
line.sort(key=lambda item: item['x'])
for item in line:
sorted_boxes.append(item['box'])
crops = []
for b in sorted_boxes:
xmin, ymin, xmax, ymax, score = b
x1 = int(xmin * W)
y1 = int(ymin * H)
x2 = int(xmax * W)
y2 = int(ymax * H)
crop = image[y1:y2, x1:x2]
h, w = crop.shape[:2]
if h == 0 or w == 0:
continue
crops.append(crop)
return crops
def predict(model, tokenizer, image_paths, wpb=16, function="generate_fast", doctr_model=None, bos_id=None, eos_id=None, device="cuda",
# generate() params
temp=None, max_iter=None, penalty=None, top_k=None,
# generate_fast() params
fast_max_iter=None,
# beam_search() params
beam_width=None, beam_max_iter=None):
"""
The main prediction function of the library.
Takes images and returns the text read from them.
Inference parameters (temp, max_iter, penalty, top_k, fast_max_iter, beam_width, beam_max_iter)
can be set here directly, OR they will be read from model.sx_config if you created the model via sx.init(config=...).
Args:
model: SOCRATE model instance.
tokenizer: SocrateXTokenizer instance.
image_paths (str | List[str]): Path(s) to the image(s).
wpb (int): Words per batch. Default: 16.
function (str | callable): 'generate', 'generate_fast', 'beam_search', or a custom callable.
doctr_model: Pre-loaded doctr model (avoids re-loading on each call).
bos_id (int): Override BOS token ID.
eos_id (int): Override EOS token ID.
device (str): Target device. Default: "cuda".
temp (float): Temperature for generate(). Default: from config or 0.5.
max_iter (int): Max tokens for generate(). Default: from config or 64.
penalty (float): Repetition penalty for generate(). Default: from config or 1.15.
top_k (int): Top-k for generate(). Default: from config or 5.
fast_max_iter (int): Max tokens for generate_fast(). Default: from config or 32.
beam_width (int): Number of beams for beam_search(). Default: from config or 4.
beam_max_iter (int): Max tokens for beam_search(). Default: from config or 64.
"""
model.eval()
# Pull inference defaults from sx_config if they were set
sx_cfg = getattr(model, "sx_config", None)
_temp = temp if temp is not None else (sx_cfg.temp if sx_cfg else 0.5)
_max_iter = max_iter if max_iter is not None else (sx_cfg.max_iter if sx_cfg else 64)
_penalty = penalty if penalty is not None else (sx_cfg.penalty if sx_cfg else 1.15)
_top_k = top_k if top_k is not None else (sx_cfg.top_k if sx_cfg else 5)
_fast_max = fast_max_iter if fast_max_iter is not None else (sx_cfg.fast_max_iter if sx_cfg else 32)
_beam_width = beam_width if beam_width is not None else (sx_cfg.beam_width if sx_cfg else 4)
_beam_max = beam_max_iter if beam_max_iter is not None else (sx_cfg.beam_max_iter if sx_cfg else 64)
# Resolve tokens (default to tokenizer if not provided)
if bos_id is None:
bos_id = tokenizer.token_to_id("<bos>")
if eos_id is None:
eos_id = tokenizer.token_to_id("<eos>")
results = {}
if isinstance(image_paths, str):
image_paths = [image_paths]
for image_path in image_paths:
crops = extract_crops_from_image(image_path, doctr_model)
if not crops:
results[image_path] = ""
continue
dataset = Makeset(images=crops)
dataloader = DataLoader(dataset, batch_size=wpb, shuffle=False, collate_fn=dataset.collate_fn)
doc_text = []
for batch in dataloader:
batch = batch.to(device)
for img in batch:
img = img.unsqueeze(0) # [1, C, H, W]
if function == "generate":
pred_ids = generate(model, img, bos_id=bos_id, eos_id=eos_id, device=device,
temp=_temp, max_iter=_max_iter, penalty=_penalty, top_k=_top_k)
elif function == "generate_fast":
pred_ids = generate_fast(model, img, bos_id=bos_id, eos_id=eos_id, device=device,
max_iter=_fast_max)
elif function == "beam_search":
pred_ids = beam_search(model, img, bos_id=bos_id, eos_id=eos_id, device=device,
beam_width=_beam_width, max_iter=_beam_max)
else:
if callable(function):
pred_ids = function(model, img, bos_id, eos_id, device)
else:
raise ValueError(f"Unknown function: {function}")
text = tokenizer.decode(pred_ids)
doc_text.append(text)
results[image_path] = " ".join(doc_text)
return results
|