digitizer / app.py
HamidAmiri's picture
Create app.py
35aaa19 verified
Raw
History Blame Contribute Delete
980 Bytes
import gradio as gr
from PIL import Image
from transformers import Pix2StructProcessor, Pix2StructForConditionalGeneration
# SMART RUNTIME LINK: Change 'your-username/my-chart-weights' to your exact HF paths
model_repo_path = "your-username/my-chart-weights"
print("Downloading and caching model weights from model repo...")
processor = Pix2StructProcessor.from_pretrained(model_repo_path)
model = Pix2StructForConditionalGeneration.from_pretrained(model_repo_path)
print("? Model loaded into 16GB RAM successfully!")
def extract_chart_data(img):
image = Image.fromarray(img).convert("RGB")
inputs = processor(images=image, text="Generate underlying data table of the figure below:", return_tensors="pt")
generated_ids = model.generate(**inputs, max_new_tokens=512)
extracted_table = processor.decode(generated_ids[0], skip_special_tokens=True)
return extracted_table
demo = gr.Interface(fn=extract_chart_data, inputs="image", outputs="text")
demo.launch()