import gradio as gr
import numpy
import pathlib
import tensorflow as tf
from PIL import Image
from music21 import stream,chord,duration,clef

myModel = tf.keras.models.load_model('my_model.h5')
class_names = ['Crash',
 'Crash + KickBass',
 'Crash + Snare',
 'Crash + Snare + KickBass',
 'Crash + Tom',
 'Crash + Tom + KickBass',
 'Hihat',
 'Hihat + KickBass',
 'Hihat + Snare',
 'Hihat + Snare + KickBass',
 'Hihat + Tom',
 'Hihat + Tom + KickBass',
 'KickBass',
 'Ride',
 'Ride + KickBass',
 'Ride + Snare',
 'Ride + Snare + KickBass',
 'Ride + Tom',
 'Ride + Tom + KickBass',
 'Snare',
 'Snare + KickBass',
 'Snare + Tom',
 'Snare + Tom + KickBass',
 'Tom',
 'Tom + KickBass']

def PredictionToNotation(predict):
    score = stream.Score()
    drum_part = stream.Part()
    drum_part.id = 'Drums'
    for prediction in predict:
        n=chord.Chord(PredictionToChord(prediction))
        n.duration = duration.Duration(0.5)
        drum_part.append(n)
    drum_part.clef = clef.PercussionClef()
    score.insert(0, drum_part)
    score.write('lily.png',"notation")
    
    return Image.open('notation.png')
    

def PredictionToChord(prediction):
    noteToChord = []
    drumComponent = {"Hihat":"A4","Snare":"D4","KickBass":"G3","Crash":"B4","Tom":"E4","Ride":"G4"}
    for component in prediction.split(" "):
        if(component in drumComponent):
            noteToChord.append(drumComponent[component])
    return noteToChord
    
 
    
def transcribe(AudioFile):
    
    return PredictionToNotation(["Hihat + Snare"])

gr.Interface(
    fn=transcribe, 
    inputs=gr.Audio(type="filepath"), 
    outputs="image").launch()