Spaces:
Sleeping
Sleeping
# import gradio as gr | |
# from transformers import AutoProcessor, AutoModelForImageTextToText | |
# from PIL import Image | |
# import re | |
# # Load SmolDocling model & processor once | |
# processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview") | |
# model = AutoModelForImageTextToText.from_pretrained("ds4sd/SmolDocling-256M-preview") | |
# def extract_fcel_values_from_image(image, prompt_text): | |
# """Run SmolDocling on an image and return numeric values inside <fcel> tags.""" | |
# # Prepare prompt for the model | |
# messages = [ | |
# {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt_text}]} | |
# ] | |
# prompt = processor.apply_chat_template(messages, add_generation_prompt=True) | |
# inputs = processor(text=prompt, images=[image], return_tensors="pt") | |
# # Generate output | |
# outputs = model.generate(**inputs, max_new_tokens=2048) | |
# prompt_length = inputs.input_ids.shape[1] | |
# generated = outputs[:, prompt_length:] | |
# result = processor.batch_decode(generated, skip_special_tokens=False)[0] | |
# clean_text = result.replace("<end_of_utterance>", "").strip() | |
# # Extract only <fcel> values | |
# values = re.findall(r"<fcel>([\d.]+)", clean_text) | |
# values = [float(v) for v in values] # convert to floats | |
# return values, clean_text | |
# def compare_images(image1, image2, prompt_text): | |
# # Extract fcel values from both images | |
# values1, raw1 = extract_fcel_values_from_image(image1, prompt_text) | |
# values2, raw2 = extract_fcel_values_from_image(image2, prompt_text) | |
# # Calculate accuracy | |
# if len(values1) == len(values2) and values1 == values2: | |
# accuracy = 100.0 | |
# else: | |
# matches = sum(1 for a, b in zip(values1, values2) if a == b) | |
# total = max(len(values1), len(values2)) | |
# accuracy = (matches / total) * 100 if total > 0 else 0 | |
# return { | |
# # "Extracted Values 1": values1, | |
# # "Extracted Values 2": values2, | |
# "Accuracy (%)": accuracy | |
# } | |
# # Gradio UI | |
# demo = gr.Interface( | |
# fn=compare_images, | |
# inputs=[ | |
# gr.Image(type="pil", label="Upload First Table Image"), | |
# gr.Image(type="pil", label="Upload Second Table Image"), | |
# gr.Textbox(lines=1, placeholder="Enter prompt (e.g. Extract table as OTSL)", label="Prompt") | |
# ], | |
# outputs="json", | |
# title="Table Data Accuracy Checker (SmolDocling)", | |
# description="Uploads two table images, extracts only <fcel> values from OTSL output, and compares them for accuracy." | |
# ) | |
# demo.launch() | |
import gradio as gr | |
from transformers import AutoProcessor, AutoModelForVision2Seq, TextIteratorStreamer | |
from transformers.image_utils import load_image | |
from threading import Thread | |
import torch | |
import html | |
import re | |
from PIL import Image, ImageOps | |
# Load model & processor once at startup | |
processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview") | |
model = AutoModelForVision2Seq.from_pretrained("ds4sd/SmolDocling-256M-preview", torch_dtype=torch.bfloat16).to("cuda") | |
def add_random_padding(image, min_percent=0.1, max_percent=0.10): | |
image = image.convert("RGB") | |
width, height = image.size | |
pad_w_percent = random.uniform(min_percent, max_percent) | |
pad_h_percent = random.uniform(min_percent, max_percent) | |
pad_w = int(width * pad_w_percent) | |
pad_h = int(height * pad_h_percent) | |
corner_pixel = image.getpixel((0, 0)) # Top-left corner | |
padded_image = ImageOps.expand(image, border=(pad_w, pad_h, pad_w, pad_h), fill=corner_pixel) | |
return padded_image | |
def extract_table(image_file): | |
# Load image | |
image = load_image(image_file) | |
# Optionally add padding if needed for model robustness (optional) | |
image = add_random_padding(image) | |
# Fixed prompt to extract table only (modify if needed) | |
text = "Convert this table to OTSL." | |
# Build the message structure for processor | |
resulting_messages = [{ | |
"role": "user", | |
"content": [{"type": "image"}] + [{"type": "text", "text": text}] | |
}] | |
prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True) | |
inputs = processor(text=prompt, images=[image], return_tensors="pt").to('cuda') | |
generation_args = { | |
"input_ids": inputs.input_ids, | |
"pixel_values": inputs.pixel_values, | |
"attention_mask": inputs.attention_mask, | |
"max_new_tokens": 8192, | |
"num_return_sequences": 1, | |
} | |
streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True) | |
generation_args["streamer"] = streamer | |
thread = Thread(target=model.generate, kwargs=generation_args) | |
thread.start() | |
output_text = "" | |
for new_text in streamer: | |
output_text += new_text | |
# Clean and return output | |
cleaned_output = output_text.replace("<end_of_utterance>", "").strip() | |
# Optionally convert <chart> tags to <otsl> if present | |
if "<chart>" in cleaned_output: | |
cleaned_output = cleaned_output.replace("<chart>", "<otsl>").replace("</chart>", "</otsl>") | |
cleaned_output = re.sub(r'(<loc_500>)(?!.*<loc_500>)<[^>]+>', r'\1', cleaned_output) | |
return cleaned_output or "No table found or unable to extract." | |
# Gradio UI | |
demo = gr.Interface( | |
fn=extract_table, | |
inputs=gr.Image(type="filepath", label="Upload Table Image"), | |
outputs=gr.Textbox(label="Extracted Table (OTSL Format)"), | |
title="Table Extraction from Image using SmolDocling-256M", | |
description="Upload an image containing a table. The model will extract the table and output it in OTSL format." | |
) | |
demo.launch(debug=True) | |