import gradio as gr
import os
from core.predict import ImageClassifier
from PIL import Image

cwd = os.getcwd()
model_path = os.path.join(cwd, "model", "cnn_model.pth")

classifier = ImageClassifier(model_path=model_path, class_names={0: 'person', 1: 'Dog', 2: 'Cat'})


def classify_image(image):
    image_path = "uploaded_image.jpg"
    image.save(image_path)
    
    label, output_path = classifier.predict(image_path)
    
    return label, Image.open(output_path)
    
demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Textbox(label = "Prediction"), gr.Image(label = "Labelled Image")],
    title = " Image Classification app ",
    description = " Upload an image to classify as Cat, Dog or Person "

)

if __name__ == "__main__":
    demo.launch()