| import gradio as gr |
| from PIL import Image |
| from transformers import Pix2StructProcessor, Pix2StructForConditionalGeneration |
|
|
| |
| 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() |