import gradio as gr
from PIL import Image
#import tensorflow as tf
#from tensorflow.keras.models import load_model
#import numpy as np

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("image-to-text", model="Flova/omr_transformer")

# # Load model directly
# from transformers import AutoTokenizer, AutoModel

# tokenizer = AutoTokenizer.from_pretrained("Flova/omr_transformer")
# model = AutoModel.from_pretrained("Flova/omr_transformer")

# Using Flova/omr_transformer

def notation_2_note(input_img):
    prediction = pipe(Image.fromarray(input_img))
    #print(type(prediction))
    output_text = prediction[0]['generated_text']
    return output_text

demo = gr.Interface(notation_2_note, gr.Image(), "text")
demo.launch()