TLME commited on
Commit
7b9fe37
1 Parent(s): f625f7d

Upload infer.py

Browse files
Files changed (1) hide show
  1. infer.py +27 -0
infer.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+
3
+ import os
4
+ import json
5
+ from mmpretrain import ImageClassificationInferencer
6
+
7
+
8
+ path = './testimg/'
9
+ config = 'convnext-v2-tiny_32xb32_in1k-384px.py'
10
+ checkpoint = 'ConvNeXt_v2-v2_ep90.pth'
11
+
12
+
13
+ inferencer = ImageClassificationInferencer(model=config, pretrained=checkpoint, device='cuda')
14
+
15
+ result={}
16
+
17
+ for root, dirs, files in os.walk(path):
18
+ for file in files:
19
+ if file.lower().endswith(('.png', '.jpg','jpeg')):
20
+ # print(os.path.join(root, file))
21
+ inf_result = inferencer(os.path.join(root, file))[0]
22
+ # print(result['pred_class'])
23
+ print(result,os.path.join(root, file))
24
+ result[os.path.join(root, file)]= [{'pred_class' : inf_result['pred_class']},{'pred_score' : inf_result['pred_score']}]
25
+
26
+ with open(path + "predict_result.json", "w") as file:
27
+ json.dump(result, file, ensure_ascii=False,indent=2)