Skip to content

Commit 0c9a78b

Browse files
quedcodeflash-ai[bot]aseembits93
authored
enhancement: optimize outputs to objects metbo2xp (#443)
From codeflash: Optimizes table model pipeline by keeping the data as a tensor for longer, enabling vector operations and reducing conversions. --------- Co-authored-by: codeflash-ai[bot] <148906541+codeflash-ai[bot]@users.noreply.github.com> Co-authored-by: aseembits93 <[email protected]>
1 parent 18c73ca commit 0c9a78b

File tree

3 files changed

+13
-7
lines changed

3 files changed

+13
-7
lines changed

CHANGELOG.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
## 1.0.8-dev0
2+
3+
* Enhancement: Optimized `outputs_to_objects` for an 88% speedup in some cases (codeflash)
4+
15
## 1.0.7
26

37
* Fix a hardcoded file extension causing confusion in the logs
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
__version__ = "1.0.7" # pragma: no cover
1+
__version__ = "1.0.8-dev0" # pragma: no cover

unstructured_inference/models/tables.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -213,16 +213,18 @@ def outputs_to_objects(
213213
):
214214
"""Output table element types."""
215215
m = outputs["logits"].softmax(-1).max(-1)
216-
pred_labels = list(m.indices.detach().cpu().numpy())[0]
217-
pred_scores = list(m.values.detach().cpu().numpy())[0]
216+
pred_labels = m.indices.detach().cpu().numpy()[0]
217+
pred_scores = m.values.detach().cpu().numpy()[0]
218218
pred_bboxes = outputs["pred_boxes"].detach().cpu()[0]
219219

220220
pad = outputs.get("pad_for_structure_detection", 0)
221221
scale_size = (img_size[0] + pad * 2, img_size[1] + pad * 2)
222-
pred_bboxes = [elem.tolist() for elem in rescale_bboxes(pred_bboxes, scale_size)]
222+
rescaled = rescale_bboxes(pred_bboxes, scale_size)
223223
# unshift the padding; padding effectively shifted the bounding boxes of structures in the
224224
# original image with half of the total pad
225-
shift_size = pad
225+
if pad != 0:
226+
rescaled = rescaled - pad
227+
pred_bboxes = rescaled.tolist()
226228

227229
objects = []
228230
for label, score, bbox in zip(pred_labels, pred_scores, pred_bboxes):
@@ -232,7 +234,7 @@ def outputs_to_objects(
232234
{
233235
"label": class_label,
234236
"score": float(score),
235-
"bbox": [float(elem) - shift_size for elem in bbox],
237+
"bbox": bbox,
236238
},
237239
)
238240

@@ -279,7 +281,7 @@ def rescale_bboxes(out_bbox, size):
279281
"""Rescale relative bounding box to box of size given by size."""
280282
img_w, img_h = size
281283
b = box_cxcywh_to_xyxy(out_bbox)
282-
b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
284+
b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32, device=out_bbox.device)
283285
return b
284286

285287

0 commit comments

Comments
 (0)