taprosoft
commited on
Commit
·
22e6cb5
1
Parent(s):
73d7517
fix: resize image
Browse files- backends/smoldocling.py +6 -2
backends/smoldocling.py
CHANGED
|
@@ -3,7 +3,6 @@
|
|
| 3 |
# pip install docling_core
|
| 4 |
# pip install transformers
|
| 5 |
|
| 6 |
-
import spaces
|
| 7 |
import torch
|
| 8 |
from docling_core.types.doc import DoclingDocument
|
| 9 |
from docling_core.types.doc.document import DocTagsDocument
|
|
@@ -34,7 +33,6 @@ messages = [
|
|
| 34 |
]
|
| 35 |
|
| 36 |
|
| 37 |
-
@spaces.GPU(duration=120)
|
| 38 |
def convert_smoldocling(path: str, file_name: str):
|
| 39 |
doc = PDF(path)
|
| 40 |
output_md = ""
|
|
@@ -42,6 +40,12 @@ def convert_smoldocling(path: str, file_name: str):
|
|
| 42 |
for image in doc.images[:MAX_PAGES]:
|
| 43 |
# convert ndarray to Image
|
| 44 |
image = Image.fromarray(image)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 45 |
|
| 46 |
# Prepare inputs
|
| 47 |
prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
|
|
|
|
| 3 |
# pip install docling_core
|
| 4 |
# pip install transformers
|
| 5 |
|
|
|
|
| 6 |
import torch
|
| 7 |
from docling_core.types.doc import DoclingDocument
|
| 8 |
from docling_core.types.doc.document import DocTagsDocument
|
|
|
|
| 33 |
]
|
| 34 |
|
| 35 |
|
|
|
|
| 36 |
def convert_smoldocling(path: str, file_name: str):
|
| 37 |
doc = PDF(path)
|
| 38 |
output_md = ""
|
|
|
|
| 40 |
for image in doc.images[:MAX_PAGES]:
|
| 41 |
# convert ndarray to Image
|
| 42 |
image = Image.fromarray(image)
|
| 43 |
+
# resize image to maximum width of 1200
|
| 44 |
+
max_width = 1200
|
| 45 |
+
if image.width > max_width:
|
| 46 |
+
image = image.resize(
|
| 47 |
+
(max_width, int(max_width * image.height / image.width))
|
| 48 |
+
)
|
| 49 |
|
| 50 |
# Prepare inputs
|
| 51 |
prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
|