|
|
|
|
|
import os |
|
import json |
|
from mmpretrain import ImageClassificationInferencer |
|
|
|
|
|
path = './testimg/' |
|
config = 'convnext-v2-tiny_32xb32_in1k-384px.py' |
|
checkpoint = 'ConvNeXt_v2-v2_ep90.pth' |
|
|
|
|
|
inferencer = ImageClassificationInferencer(model=config, pretrained=checkpoint, device='cuda') |
|
|
|
result={} |
|
|
|
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)) |
|
result[os.path.join(root, file)]= [{'pred_class' : inf_result['pred_class']},{'pred_score' : inf_result['pred_score']}] |
|
|
|
with open(path + "predict_result.json", "w") as file: |
|
json.dump(result, file, ensure_ascii=False,indent=2) |
