File size: 2,580 Bytes
6000943 |
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 |
import os
import json
import pandas as pd
class DAGMSolver(object):
CLSNAMES = [
'Class1','Class2','Class3','Class4','Class5','Class6','Class7','Class8','Class9','Class10'
]
def __init__(self, root='data/mvtec'):
self.root = root
self.meta_path = f'{root}/meta.json'
def run(self):
info = dict(train={}, test={})
anomaly_samples = 0
normal_samples = 0
for cls_name in self.CLSNAMES:
cls_dir = f'{self.root}/{cls_name}'
for phase in ['train', 'test']:
cls_info = []
x, y, mask_names_none= [], [], []
img_dir = os.listdir(f'{cls_dir}/{phase}')
mask_names = os.listdir(f'{cls_dir}/{phase}/Label')
img_fpath_list = sorted([f
for f in img_dir
if f.endswith('.PNG')])
gt_fpath_list = sorted([f
for f in mask_names
if f.endswith('.PNG')])
img_exclude_list = [f.split("_")[0] + ".PNG" for f in gt_fpath_list]
img_normal_fpath_list = list(set(img_fpath_list) - set(img_exclude_list))
x.extend(img_normal_fpath_list + img_exclude_list)
y.extend([0] * len(img_normal_fpath_list) + [1]* len(img_exclude_list))
mask_names_none.extend([None] * len(img_normal_fpath_list) + gt_fpath_list)
for idx, img_name in enumerate(x):
info_img = dict(
img_path=f'{cls_name}/{phase}/{img_name}',
mask_path=f'{cls_name}/{phase}/Label/{mask_names_none[idx]}' if mask_names_none[idx] != None else '',
cls_name=cls_name,
specie_name='',
anomaly=1 if y[idx] == 1 else 0,
)
cls_info.append(info_img)
if phase == 'test':
if y[idx] == 1:
anomaly_samples = anomaly_samples + 1
else:
normal_samples = normal_samples + 1
info[phase][cls_name] = cls_info
with open(self.meta_path, 'w') as f:
f.write(json.dumps(info, indent=4) + "\n")
print('normal_samples', normal_samples, 'anomaly_samples', anomaly_samples)
if __name__ == '__main__':
runner = DAGMSolver(root='/home/hzx/Projects/Data/DAGM/')
runner.run() |