#notes: cek checkpoint
import os
import torch
from basic import ResNet18
import gradio as gr
import numpy as np
from torchvision.transforms import transforms
from transformers import AutoModelForSequenceClassification

# Gantilah dengan nama yang Anda berikan saat mengunggah checkpoint
#model_name = 'hasanah10105/breast-cancer-classification/ckpt/epoch=49-step=1750.ckpt'

# Tentukan path ke file checkpoint Anda di Hugging Face
#checkpoint_path = "hasanah10105/ckpt/epoch=49-step=1750.ckpt"
#checkpoint_path = r"C:\Users\Asus\AI-Vision-Lightning-Pipeline\breast-cancer-classification\ckpt\epoch=49-step=1750.ckpt"

# Muat model dari Hugging Face
model = AutoModelForSequenceClassification.from_pretrained('hasanah10105/breast-cancer-classification/ckpt/epoch=49-step=1750.ckpt')

# Muat parameter model dari checkpoint Anda
#model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
#ROOT_DIR = os.path.abspath('')
#PRETRAINED_MODEL = ./ckpt/epoch=49-step=1750.ckpt
 

#model = ResNet18(3, 3)

#checkpoint = torch.load(model)

class_names = ['benign', 'malignant', 'normal']
class_names.sort()

transformation_pipeline = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Grayscale(num_output_channels=1),
        transforms.Resize((256, 256)),
        transforms.RandomRotation(20),
        transforms.ToTensor(), 
        transforms.Normalize(mean=[0], std=[1])
])

def preprocess_image(image: np.ndarray):
    """Preprocess the input image.

    Note that the input image is in RGB mode.

    Parameters
    ----------
    image: np.ndarray
        Input image from callback.
    """

    image = transformation_pipeline(image)
    image = torch.unsqueeze(image, 0)

    return image

def image_classifier(inp):
    """Image Classifier Function.

    Parameters
    ----------
    inp: Optional[np.ndarray] = None
        Input image from callback

    Returns
    -------
    Dict
        A dictionary class names and its probability
    """

    # If input not valid, return dummy data or raise error
    if inp is None:
      return {gr.Error()}

    # preprocess
    image = preprocess_image(inp)
    image = image.to(dtype=torch.float32)

    # inference
    result = model(image)

    # postprocess
    result = torch.nn.functional.softmax(result.logits, dim=1) # apply softmax
    result = result[0].detach().numpy().tolist() # take the first batch
    labeled_result = {name:score for name, score in zip(class_names, result)}

    return labeled_result

iface = gr.Interface(fn=image_classifier, inputs="image", outputs="label")
iface.launch()