File size: 3,195 Bytes
6d3c96e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
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()
|