import torch
import requests
import gradio as gr
from PIL import Image
from transformers import ResNetForImageClassification, AutoImageProcessor

target_folder = "Kang-Seong-Jun/Korean_Real_Estate_Classifier"

def load_model_and_preprocessor(target_folder):
    model = ResNetForImageClassification.from_pretrained(target_folder)
    image_processor = AutoImageProcessor.from_pretrained(target_folder)
    return model, image_processor

def infer_image(image, model, image_processor, k):
    processed_img = image_processor(images=image.convert("RGB"), return_tensors="pt")

    with torch.no_grad():
        outputs = model(**processed_img)
        logits = outputs.logits
    
    prob = torch.nn.functional.softmax(logits, dim=-1)
    topk_prob, topk_indices = torch.topk(prob, k=k)   
    
    res = ""
    for idx, (prob, index) in enumerate(zip(topk_prob[0], topk_indices[0])):
        res += f"{idx+1}. {model.config.id2label[index.item()]:<15} ({prob.item()*100:.2f} %) \n" 
    return res

def infer(image, k, target_folder=target_folder):
    try: 
        model, image_processor = load_model_and_preprocessor(target_folder)
        res = infer_image(image, model, image_processor, k)
    except Exception as e: 
        image = Image.new('RGB', (224, 224))
        res = f"이미지를 처리하는데 문제가 발생했습니다: {str(e)}"
    return image, res 

demo = gr.Interface(
    fn=infer,
    inputs=[
        gr.Image(type="pil", label="입력 이미지"),
        gr.Slider(minimum=0, maximum=20, step=1, value=3, label="상위 몇개까지 보여줄까요?")
    ],
    outputs=[
        gr.Image(type="pil", label="입력 이미지"),
        gr.Textbox(label="종류 (확률)")
    ],
)

demo.launch()