import os
import json
from mmpretrain import ImageClassificationInferencer
import torch
import gradio as gr

config = 'convnext-v2-tiny_32xb32_in1k-384px.py'
checkpoint = 'ConvNeXt_v2-v2_ep90.pth'


inferencer = ImageClassificationInferencer(model=config, pretrained=checkpoint, device= "cuda" if torch.cuda.is_available() else "cpu")

def single_image_classifier(image):
    inf_result = inferencer(image)[0]
    label = inf_result['pred_class']
    score = inf_result['pred_score']
    if label == "not_western":
        another_label = "western"
        another_score = 1 - score
    else:
        another_label = "not_western"
        another_score = 1 - score
    return {label : score,another_label:another_score}

def batch_process(path,is_pred_score):
    result={}
    try:
        for root, dirs, files in os.walk(path):
            for file in files:
                if file.lower().endswith(('.png', '.jpg','jpeg')):
                    inf_result = inferencer(os.path.join(root, file))[0]
                    print(result,os.path.join(root, file))
                    if is_pred_score == True:
                        result[os.path.join(root, file)]= [{'pred_class' : inf_result['pred_class']},{'pred_score' : inf_result['pred_score']}]
                    else:
                        result[os.path.join(root, file)]= [{'pred_class' : inf_result['pred_class']}]
        with open(path+ "/" + "predict_result.json", "w") as file:
            json.dump(result, file, ensure_ascii=False,indent=2)
        return "sucess"
    except:
        return "failed"

with gr.Blocks() as demo:
    gr.Markdown("# Western anime images classification")
    gr.Markdown("A classification using mmpretrain trained to classify western images based on ConvNeXtV2-tiny.Used for classifying anime images based on whether they are in the Western style.\n\n"
                "The inference script: https://huggingface.co/TLME/western-classification"
                )
    with gr.Tab("Single image"):
            input_img = gr.Image(source='upload')
            
            
            output_label =gr.Label(label="Predict result")
            examples_imgs = ["./testimg2/1.jpg","./testimg2/2.jpg","./testimg2/3.jpg","./testimg2/4.jpg","./testimg/1.jpg","./testimg/2.jpg","./testimg/3.jpg","./testimg/4.jpg"]
            
            button = gr.Button("Submit",variant="primary")
            
            button.click(single_image_classifier,inputs= input_img,outputs= output_label)
            gr.Examples(examples= examples_imgs ,inputs = input_img , outputs= output_label,fn=single_image_classifier)
            
    with gr.Tab("Batch process"):
        with gr.Row(label ='Input path'):
            with gr.Column():
                input_path = gr.Textbox(label="input your images folder")
                is_pred_score = gr.Checkbox(value = True, label="Output pred_score")
            
        output_msg = gr.Textbox(label="Message")

        buttom =gr.Button("Process",variant="primary")
        buttom.click(batch_process, inputs= [input_path,is_pred_score] , outputs=output_msg)
        
if __name__ == "__main__":
    demo.queue(concurrency_count=4)
    demo.launch()