import numpy as np
import gradio as gr
from PIL import Image
import torch
import os
import pytorch_lightning as pl

MODEL_PATH = "./model.pth"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = None
try:
    model = torch.load(MODEL_PATH, weights_only=False)
    model = model.to(device)
    print("Model Loaded Successfully")
except Exception as e:
    print(e)

def process_image(image):
    image = image.convert("L") # Converts into grayscale image
    image = image.resize((28, 28)) # resizes into shapes that was in training 
    image = np.array(image) / 255.0 # pixels normalizes into [0, 1]
    image = (image - 0.1307) / 0.3081 # standard normalization
    image = torch.tensor(image) # converts the image from np to torch 1x28x28
    image = image.unsqueeze(dim=0) # adds a batch dimension 1, 1, 28, 28
    return image.to(device)

def predict_image(image_path):
    image = Image.open(image_path) # reads the image as PIL image
    image = process_image(image)
    
    image = image.float()   
    image = image.to(next(model.parameters()).device)
    
    try:
        model.eval() # set the mode as evaluation
        with torch.no_grad():
            output = model(image) # outputs (1, 10) [0.2, 0.1, 0.05, 0., 0., 0., 0., 0., 0., 0.6, 0.05]
            prediction = output.argmax(dim=1)
            prediction = prediction.item()
            return f"The digit is {prediction}"
    except Exception as e:
        return str(e)

interface = gr.Interface(
    fn=predict_image,
    inputs=gr.components.Image(type='filepath'),
    outputs=gr.components.Label(),
    title="Hand Written Digit Recognition App",
    description="Upload a grayscale image."
)

interface.launch(share=False)