Update main.py

main.py

@@ -6,6 +6,7 @@ import subprocess
 from torchvision import transforms
 from PIL import Image
 
+
 DEFAULT_FILE_NAME = 'sheet_music'
 
 
@@ -38,6 +39,8 @@ def convert_to_key():
     return f'{file_name}.musicxml', f'{file_name}_teaser.png'
 
 def predict(inp):
+    img = Image.fromarray(inp)
+    img.save(f'{DEFAULT_FILE_NAME}.png')
     return convert_to_key()
 
 
@@ -46,7 +49,7 @@ def run():
                gr.File(label='Download detection image', type='file')]
     demo = gr.Interface(
         fn=predict,
-        inputs='text',
+        inputs='image',
         outputs=outputs,
         allow_flagging='never',
     )