model = tf.keras.models.load_model('mnist.keras') # .keras modellunk betoltese

def preprocess_image(img):  # Szokasos preprocess eljaras
    # Grayscale konverzio (ha nem ugy lenne)
    if len(img.shape) == 3 and img.shape[-1] == 3:
        img = tf.image.rgb_to_grayscale(img).numpy().squeeze()
    
    img = tf.image.resize(img[None, ..., None], (28, 28)).numpy().squeeze() # atmeretez: 28x28
    
    # szures: fekete hatter - feher digitek
    if img.mean() > 127:
        img = 255 - img
    
    img = img.astype('float32') / 255.0 # Normalize
    img = img.reshape(784) # Lapositjuk az adatot mert az ANN igy keri
    
    return img

def predict_digit(img):  # Predict funkcio - Ez indul a submit gombra
    processed_img = preprocess_image(img)  # Preprocess the image
    processed_img = np.expand_dims(processed_img, axis=0)   # Add batch dimension
    prediction = model.predict(processed_img)[0]  # Make prediction
    
    # Return the prediction as a dictionary
    return {str(i): float(prediction[i]) for i in range(10)}

# A gradio interface
iface = gr.Interface(
    fn=predict_digit,   # a fo funkcio
    inputs=gr.Image(type="numpy", label="Tölts fel egy képet ..."),
    outputs=gr.Label(num_top_classes=3, label="A feltöltött számjegy:"),
    title="MNIST Digit Recognition",
    description="Tölts fel egy kézzel írott szám képet (0-9) és a modell megállapítja melyik szám az",
)

# iface.launch() # Interface indítása
iface.launch(share=True) # Interface indítása