| 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() | |