| import os
|
| import json
|
|
|
|
|
| class DATASolver(object):
|
|
|
| CLSNAMES = [
|
| 'shinpyung',
|
|
|
| ]
|
|
|
| def __init__(self, root='data/4inlab'):
|
| self.root = root
|
| self.meta_path = f'{root}/meta_train.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 = []
|
| species = os.listdir(f'{cls_dir}/{phase}')
|
| for specie in species:
|
| is_abnormal = True if specie not in ['good'] else False
|
| img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
|
| mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
|
| img_names.sort()
|
| mask_names.sort() if mask_names is not None else None
|
| for idx, img_name in enumerate(img_names):
|
| info_img = dict(
|
| img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
|
| mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
|
| cls_name=cls_name,
|
| specie_name=specie,
|
| anomaly=1 if is_abnormal else 0,
|
| )
|
| cls_info.append(info_img)
|
| if phase == 'test':
|
| if is_abnormal:
|
| 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 = DATASolver(root='data/4inlab')
|
| runner.run()
|
|
|
