from fastai.vision.all import *
from fastai.learner import load_learner
import gradio as gr
import numpy as np
from PIL import Image
import pickle

def get_label(file_path):
    image_id = file_path.stem
    matching_rows = df.loc[df['image_id'] == image_id, 'is_cancerous']

    if matching_rows.empty:
        # print(f"Warning: No matching row found for image_id: {image_id}")
        return 0

    # if len(matching_rows) > 1:
    #     print(f"Warning: Multiple rows found for image_id: {image_id}. Using the first one.")

    label = matching_rows.iloc[0]
    # print(f"Label for {image_id}: {label}")
    return label

learn = load_learner('skin_lesion_model.pkl')

labels = learn.dls.vocab

label_map = {0: "Non-cancerous", 1: "Cancerous"}

def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learn.predict(img)
    return {label_map[i]: float(probs[i]) for i in range(len(labels))}

gr.Interface(
    fn=predict, 
    inputs=gr.Image(), 
    outputs=gr.Label(num_top_classes=2),
    title="Skin Lesion Classification",
    description="Upload an image to classify it as cancerous or non-cancerous."
).launch()