Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import AutoProcessor, Pix2StructForConditionalGeneration | |
| import torch | |
| from PIL import Image | |
| import json | |
| import vl_convert as vlc | |
| from io import BytesIO | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Load the processor and model | |
| processor = AutoProcessor.from_pretrained("google/matcha-base") | |
| processor.image_processor.is_vqa = False | |
| model = Pix2StructForConditionalGeneration.from_pretrained("martinsinnona/visdecode_B").to(device) | |
| model.eval() | |
| def generate(image): | |
| #inputs = processor(images=image, return_tensors="pt", max_patches=1024).to(device) | |
| #generated_ids = model.generate(flattened_patches=inputs.flattened_patches, attention_mask=inputs.attention_mask, max_length=600) | |
| #generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
| generated_caption = "{'mark': 'bar', 'encoding': {'x': {'field': '', 'type': 'ordinal'}, 'y': {'field': '', 'type': 'quantitative'}}, 'data': {'values': [{'x': 0, 'y': 5.6}, {'x': 1, 'y': 6.7}, {'x': 2, 'y': 5.0}, {'x': 3, 'y': 18.7}]}}" | |
| # Generate the Vega image | |
| vega = string_to_vega(generated_caption) | |
| vega_image = draw_vega(vega) | |
| return generated_caption, vega_image | |
| def draw_vega(vega, scale=3): | |
| spec = json.dumps(vega, indent=4) | |
| png_data = vlc.vegalite_to_png(vl_spec=spec, scale=scale) | |
| return Image.open(BytesIO(png_data)).thumbnail((500,500)) | |
| def string_to_vega(string): | |
| string = string.replace("'", "\"") | |
| vega = json.loads(string) | |
| for axis in ["x", "y"]: | |
| field = vega["encoding"][axis]["field"] | |
| if field == "": | |
| vega["encoding"][axis]["field"] = axis | |
| vega["encoding"][axis]["title"] = "" | |
| else: | |
| for entry in vega["data"]["values"]: | |
| entry[field] = entry.pop(axis) | |
| return vega | |
| # Create the Gradio interface | |
| iface = gr.Interface( | |
| fn=generate, | |
| inputs=gr.Image(type="pil"), | |
| outputs=[gr.Textbox(), | |
| gr.Image(type="pil", label="Generated Vega Image", height=500, width=500)], | |
| title="Image to Vega-Lite", | |
| description="Upload an image to generate vega-lite" | |
| ) | |
| # Launch the interface | |
| if __name__ == "__main__": | |
| iface.launch(share=True) | |